feat(models): Florence-2 BART seq2seq engine and text core by inureyes · Pull Request #1060 · lablup/mlxcel

inureyes · 2026-08-07T03:31:49Z

Summary

Lands the reusable BART-style seq2seq engine plus the Florence-2 text core (sub-issue of epic #850): bidirectional encoder, causal decoder with encoder cross-attention, and the dual KV cache the decode loop needs. This is the foundation the vision-fusion (#854), processor (#855), and capstone (#856) sub-issues drive the real model through.

What changed

src/models/florence2/layers.rs: BART post-norm building blocks. Florence2Attention (bidirectional / causal / cross variants over one projection set), Florence2EncoderLayer, Florence2DecoderLayer, the KvCache pair, and an offset-aware additive causal mask. Post-norm ordering (x = LN(x + Sublayer(x))) per the mlx-vlm reference, unlike Whisper's pre-norm blocks.
src/models/florence2/encoder.rs: bidirectional encoder with learned absolute positions at the BART table offset of 2 and an embedding LayerNorm. Consumes pre-computed embeddings so Blaizzy/mlx-vlm-style fused image + text sequences enter the same path in feat(models): Florence-2 vision-language fusion + full weight loading (sub of #850) #854.
src/models/florence2/decoder.rs: causal decoder indexed by the running decode offset, cross-attending to the encoder output through per-layer dual caches.
src/models/florence2/mod.rs: Florence2TextConfig (parses the text_config sub-object of a real Florence-2 config.json, or a bare text config), shift_tokens_right with decoder_start_token_id = 2, Florence2SeqCache, and Florence2TextModel with the public engine API: load / from_weights (caller-supplied key prefix so the full-VLM loader reuses it), embed_tokens, encode_tokens, encode_embeds, decode, make_cache, generate_greedy. Weight loading filters to language_model.* tensors and applies the standard bf16 -> f16 conversion; a materialized lm_head.weight is preferred with a tied-shared-embedding fallback.
src/models/florence2/florence2_tests.rs: 13 unit tests on a tiny synthetic model covering config parsing, shift semantics, encoder bidirectionality, incremental-vs-full-sequence decode equality, prefill-plus-steps equality (offset causal mask), one-shot cross cache vs growing self cache, and greedy round-trip determinism.
tests/florence2_text_parity.rs: real-checkpoint parity against the mlx-vlm florence2 LanguageModel on models/Florence-2-base-ft-bf16 (skips when the checkpoint is absent). Pins teacher-forced per-position argmax, logit slices at the first and last positions, and the greedy continuation.
src/models/mod.rs: module registration and re-exports; src/lib/mlxcel-core/src/weights.rs: // Used by: comment update for load_weights_from_dir_filtered.

Correction to the issue text

The issue body says 8 attention heads; the real checkpoint config has encoder_attention_heads: 12 and decoder_attention_heads: 12 (the mlx-vlm TextConfig dataclass default of 8 is overridden by the checkpoint). The implementation and tests derive every shape from the real config.json and pin 12 heads explicitly.

Validation

Reference values were captured by running the mlx-vlm florence2 LanguageModel (imported from the upstream source) on the same checkpoint with weights cast bf16 -> f16 to match mlxcel's Apple Silicon precision policy. mlxcel reproduces the reference exactly on teacher-forced argmax [0, 2362, 16, 10, 2170, 13, 2] and greedy continuation [0, 2362] (EOS stop), with logit slices within 0.05 in f16.

Test plan

cargo test --profile test-fast --features metal,accelerate --lib florence2_tests:: (13 passed)
cargo test --profile test-fast --features metal,accelerate --test florence2_text_parity (1 passed, real checkpoint)
cargo clippy --profile test-fast --features metal,accelerate --lib --tests -- -D warnings (clean)
cargo fmt --check (clean)
python3 scripts/ci/check_cross_repo_refs.py (OK)

Closes #852

Add the reusable encoder-decoder inference path for Florence-2 under src/models/florence2/: a bidirectional BART encoder, a causal decoder with encoder cross-attention, and the dual KV cache (per-step self-attention history plus one-shot cross-attention K/V projected once from the encoder output) that the seq2seq decode loop needs. mlxcel's only prior encoder-decoder precedent is Whisper; this port mirrors that module's structure while switching from pre-norm to BART post-norm ordering and from audio-to-text to token-to-token. Florence2TextConfig parses the text_config sub-object of a real Florence-2 config.json (d_model 768, 6+6 layers, 12 heads on both sides, ffn 3072, vocab 51289, learned positions with the BART table offset of 2, decoder_start_token_id 2). shift_tokens_right implements the HF BART teacher-forcing shift. Weights load through load_weights_from_dir_filtered keeping only language_model.* tensors, with the standard bf16 to f16 conversion; the materialized lm_head.weight is used when present with a tied-shared-embedding fallback. The public API (embed_tokens / encode_tokens / encode_embeds / decode / make_cache / generate_greedy, plus from_weights with a caller-supplied key prefix) is the surface the vision-fusion and processor sub-issues drive the full model through; encode_embeds exists so fused image + text embeddings enter the same encoder path. Validation: 13 unit tests on a tiny synthetic model pin config parsing, shift semantics, encoder bidirectionality, and dual-cache correctness (incremental decode and prefill-plus-steps both match teacher-forced full-sequence logits; cross K/V stays pinned to the encoder length while self K/V grows per step). A real-checkpoint parity test against models/Florence-2-base-ft-bf16 matches the mlx-vlm florence2 LanguageModel reference: identical teacher-forced per-position argmax, logit slices within 0.05 in f16, and identical greedy continuation. Refs #852 (sub-issue of epic #850).

The learned position table has max_position_embeddings rows; an overlong or empty encoder prompt previously failed inside the MLX graph instead of as a clean error. generate_greedy now rejects both up front, and the encode_tokens / encode_embeds docs state the precondition for callers that embed sequences directly (the vision-fusion path bounds its own inputs). Validated with the florence2 unit suite (14 passed, including the new rejection test) and the real-checkpoint parity test. Refs #852.

## Summary Implements the Florence-2 DaViT ("Dual Attention Vision Transformer") vision backbone, the pixel-side half of the model that will feed the BART seq2seq engine landed in #1060. Four stages, each a strided `ConvEmbed` patch embedding followed by `depths[i]` blocks that alternate windowed spatial attention with channel attention, every half sandwiched between two depthwise positional convolutions. For the `Florence-2-base-ft` geometry and a 768x768 input the progression is 192x192 @ 128 -> 96x96 @ 256 -> 48x48 @ 512 -> 24x24 @ 1024, so the backbone emits `[B, 576, 1024]`. Scope is the backbone only. Fusion / projection, the processor, and end-to-end validation are separate sub-issues that depend on this one. ## What changed - `src/vision/encoders/florence2_davit.rs` (new): `Florence2VisionConfig` (parse + validate), the `sanitize` channels-last remap, `Florence2DaViT` with `from_weights` / `load` / `forward` / `forward_stages`, and `impl VisionEncoder`. - `src/vision/encoders/florence2_davit_blocks.rs` (new): `ConvEmbed`, `DepthWiseConv2d`, `Mlp`, `SpatialBlock`, `ChannelBlock`, `Block`. - `src/vision/encoders/florence2_davit_attention.rs` (new): `WindowAttention` and `ChannelAttention`. Split out of the blocks file to stay under the 500-line limit. - `src/vision/encoders/florence2_davit_tests.rs` (new): checkpoint-free unit tests for config parsing / validation, the sanitize pass and its idempotency, a synthetic-weight forward through the window pad / crop path, and the missing-weight error path. - `tests/florence2_vision_parity.rs` (new): numerical parity against the mlx-vlm reference on the real checkpoint, gated on checkpoint presence. - `src/vision/encoders/deepseekocr_sam.rs`: `window_partition` / `window_unpartition` reused rather than duplicated; their `// Used by:` comments now name Florence-2 DaViT, per `docs/code-guidelines.md`. - `src/vision/encoders/mod.rs`, `src/models/florence2/mod.rs`, `src/models/mod.rs`: module registration and re-export, so the fusion sub-issue can reach `Florence2DaViT` from the Florence-2 family directory alongside `Florence2TextModel`. ## API for the fusion sub-issue `Florence2VisionConfig::from_model_config(&Value)` (descends into `vision_config`, also accepts a bare vision-config object, mirroring `Florence2TextConfig::from_model_config`), `Florence2DaViT::from_weights(&WeightMap, &Florence2VisionConfig, prefix)` with `FLORENCE2_VISION_PREFIX = "vision_tower."`, `Florence2DaViT::load(&Path)` (config + `vision_tower.*` tensors + sanitize + bf16 -> f16 per the Apple Silicon precision policy), `forward` -> `[B, N, dim_embed[-1]]`, `forward_stages` -> per-stage `(tokens, (H, W))`, and `impl VisionEncoder`. ## Corrections to the issue text Every one of these was checked against `models/Florence-2-base-ft-bf16` (config.json and the actual safetensors keys and shapes) and against the upstream source, not inferred from the issue prose. 1. **`vision_config.model_type` is the empty string `""`, not `"davit"`.** The issue says to parse a `VisionConfig` with `model_type: davit`. Gating on that literal rejects every real Florence-2 checkpoint. Parsing here accepts `""`, `"davit"`, and an absent key (matching the upstream `VisionModel.__init__` guard, which allows `["davit", ""]`) and rejects anything else; the family is identified by the parent `model_type: florence2` plus the structural fields. Covered by a unit test that asserts all three spellings parse and a bogus one errors. 2. **There is no `hidden_size` in `vision_config`, and the backbone output is not 768.** The issue says "output image features with `hidden_size = 768`". The checkpoint has `projection_dim: 768` at both the top level and inside `vision_config`, while the last stage's `dim_embed` is `1024`, and `image_projection` is a `[1024, 768]` tensor. So the backbone emits `[B, N, 1024]`; the 1024 -> 768 projection is `image_projection` + `image_proj_norm` and belongs to the fusion sub-issue. `projection_dim` is parsed and retained on the config but never applied here. (Upstream's `VisionConfig` dataclass does carry a `hidden_size: int = 768` field, but `vision.py` never reads it.) 3. **The `window_partition` / `window_reverse` reuse pointer in the issue is wrong.** The issue says to reuse the patterns "from qwen2_5_vl and kimi_vl". Those encoders window differently. The helpers with matching semantics live only in `src/vision/encoders/deepseekocr_sam.rs`, and that is what this PR reuses (including the pad-inside-partition / crop-inside-unpartition contract, which lets the reference's explicit `mx.pad` and `x[:, :H, :W, :]` steps drop out). 4. **The issue omits four present, load-bearing config fields.** `num_groups: [4,8,16,32]`, `image_pos_embed: {type: "learned_abs_2d", max_pos_embeddings: 50}`, `visual_temporal_embedding: {type: "COSINE", max_temporal_embeddings: 100}`, and `image_feature_source: ["spatial_avg_pool","temporal_avg_pool"]`. Only `num_groups` affects this forward: it is `ChannelAttention.groups` per stage. The other three are consumed exclusively by `Florence2._encode_image` in the fusion stage, and upstream carries them on `ModelConfig` rather than `VisionConfig` while the real checkpoint puts them inside `vision_config`. All four are parsed and retained on `Florence2VisionConfig` so the fusion sub-issue reads them off one struct. 5. **The upstream `ChannelAttention` inline comment mis-states the attention shape.** It annotates `mx.matmul(q.transpose(0, 1, 3, 2), k)` as `(B, groups, N, N)`. The operands are `(B, g, C/g, N)` and `(B, g, N, C/g)`, so the product is `(B, groups, C/groups, C/groups)`: channel-by-channel, which is the whole point of channel attention. This port follows the code, and the doc comment on `ChannelAttention` records the discrepancy so nobody "corrects" it back. 6. **The `mlx-community` bf16 export already ships conv weights channels-last, so `sanitize` is a no-op on it.** `vision_tower.convs.0.proj.weight` is `[128, 7, 7, 3]` and `...dw.weight` is `[128, 3, 3, 1]` in the actual file, i.e. already `(O, kH, kW, I)`. The pass is still implemented in full (with the reference's `check_array_shape` guard on `convs.*.proj.weight` and the `shape[1] < shape[-1]` guard on `blocks.*.dw.weight`) so a raw PyTorch export converts correctly, and a unit test asserts both the transposing and the idempotent branches. 7. **`ChannelAttention` scales queries by the token count, not the head dimension.** `q = q * (float(N) ** -0.5)` where `N` is up to ~37k at stage 0. Unusual enough to look like a bug; it is not, and it is commented as deliberate. The softmax there runs in f32 and casts back, since the reduction axis is long. 8. **`DropPath` is training-only.** `drop_path_rate: 0.1` is present in the checkpoint and is parsed onto the config, but `self.training` is false at inference so the module is the identity. Implemented as a no-op with a comment saying why the config value does not reach the forward. `docs/supported-models.md` and `src/models/detection.rs` are deliberately untouched: this sub-issue does not make Florence-2 a loadable family (fusion does), so advertising it in the support table or wiring it into architecture detection now would claim a capability the binary does not have. Both belong with the fusion sub-issue. ## Validation **Input choice.** A deterministic synthetic pixel tensor, not a decoded image. Florence-2 image preprocessing is the processor sub-issue's scope, so reproducing it byte-identically in both runtimes is not possible yet, and what parity actually requires is that both sides see identical input bytes. Both sides build `x[0, c, h, w] = ((h * 768 + w) * 3 + c) % 251 / 251.0 - 0.5` in f32 (every step a single IEEE-754 operation over exactly represented integers, so the two implementations agree bit for bit) and then cast to f16. **Reference.** The mlx-vlm florence2 `VisionModel` (https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/models/florence2/vision.py) run in a scratch venv against `models/Florence-2-base-ft-bf16`, with `vision_tower.*` weights put through `VisionModel.sanitize` and cast bf16 -> f16 to match mlxcel's Apple Silicon precision policy. Captured per-stage shape, first-16 activations, and mean / std; the regeneration recipe is recorded in the test's module doc. **Tolerance.** Per-stage absolute, scaled to the activation magnitude at that stage, because both sides run f16 with different op ordering: `2e-3 / 1.2e-2 / 6e-2 / 1.2e-2` against stage peaks of `1.54 / 8.03 / 65.5 / 3.97`. Each is a small multiple of one f16 ulp at its stage's peak and at least 25x below that stage's own standard deviation. Nothing was loosened to make a failure pass; the tolerances were set after the first run already agreed. **Observed max absolute deviation** (Apple Silicon, `test-fast` profile): | Stage | Shape | Max abs deviation | Tolerance | Mean (ours / ref) | Std (ours / ref) | |-------|-------|-------------------|-----------|-------------------|------------------| | 0 | `[1, 36864, 128]` | 3.37e-4 | 2e-3 | -0.001657 / -0.001656 | 0.119353 / 0.119339 | | 1 | `[1, 9216, 256]` | 4.29e-3 | 1.2e-2 | 0.010804 / 0.010804 | 0.290855 / 0.290888 | | 2 | `[1, 2304, 512]` | 7.38e-4 | 6e-2 | 0.074498 / 0.074503 | 1.458913 / 1.458821 | | 3 | `[1, 576, 1024]` | 3.14e-3 | 1.2e-2 | -0.487351 / -0.487382 | 0.524581 / 0.524704 | ## Test plan - [x] `cargo test --profile test-fast --features metal,accelerate -p mlxcel --lib vision::encoders::florence2_davit` (7 passed) - [x] `cargo test --profile test-fast --features metal,accelerate -p mlxcel --test florence2_vision_parity -- --nocapture` (2 passed) - [x] `cargo test --profile test-fast --features metal,accelerate -p mlxcel --test florence2_text_parity` (1 passed, no regression from the re-export) - [x] `cargo clippy --profile test-fast --features metal,accelerate --lib --tests -- -D warnings` (clean) - [x] `cargo fmt --check` (clean) - [x] `python3 scripts/ci/check_cross_repo_refs.py` (clean) Closes #853

…ng (#1064) ## Summary Joins the DaViT vision tower (#1063) to the BART seq2seq engine (#1060) into `Florence2Model`, the assembled Florence-2 vision-language model, and completes weight loading for the whole checkpoint. Florence-2 concatenates rather than scattering: the tower output is pooled into a short feature sequence, projected into the text embedding space, and placed in front of the task-prompt embeddings to form the encoder input, with no image placeholder token in the prompt at all. ## What changed - `src/models/florence2/fusion.rs` (new): `LearnedPositionEmbedding2D` (learned column half first, then row half, over a square feature grid) and `PositionalEmbeddingCosine1D` (the `visual_temporal_embed.pos_idx_to_embed` table, taken from the checkpoint when present and synthesized from the closed form otherwise), plus `additive_attention_mask`, which turns a 0/1 mask into the `[batch, 1, 1, seq]` additive form attention consumes. - `src/models/florence2/checkpoint.rs` (new): `Florence2Config` (both halves plus `image_token_id`, defaulting to the text vocabulary size) and `sanitize`, which runs the DaViT conv channels-last remap over the whole weight map, drops the BART `final_logits_bias` buffer, and fills `model.shared.weight` from an `embed_tokens` copy for exports that materialize only the tied table. - `src/models/florence2/model.rs` (new): `Florence2Model` with `encode_image` / `encode_image_features` (`_encode_image`: 2-D grid position embedding, temporal embedding, `image_feature_source` pooling, `image_projection` 1024 to 768, `image_proj_norm`), `embed_prompt`, `merge_input_ids_with_image_features`, `encode` / `encode_with_image_features` / `encode_fused`, `make_cache`, `decode`, and `generate_greedy`. - `src/models/florence2/{encoder,layers}.rs`: the encoder and its blocks now take an optional additive attention mask, so the joint mask the fusion stage builds reaches the encoder's self-attention instead of being dropped. `encode_tokens` / `encode_embeds` pass `None`, leaving the text-only numerics unchanged; the new `Florence2TextModel::encode_embeds_with_mask` is what the fused path drives. - `src/models/florence2/florence2_fusion_tests.rs` (new): 37 unit tests covering both positional embeddings, the additive mask, all three `sanitize` passes, config parsing, load-time validation, and the whole `pixels -> features -> concat -> encoder -> decoder logits` chain on a tiny synthetic model (one block-free DaViT stage plus a one-layer BART stack), so the integrated path is exercised without a checkpoint. - `tests/florence2_fusion_parity.rs` (new): gated end-to-end parity against the mlx-vlm florence2 `Model` on `models/Florence-2-base-ft-bf16`, skipping cleanly when the checkpoint is absent. ## Precision and dtype handling The fused path pins its activation dtype from the text stack and casts the projected image features to it, and casts pixel input to the vision tower's own dtype, so a checkpoint whose two halves land on different dense dtypes does not silently promote the fused sequence to f32. Weight loading applies the usual Apple Silicon policy of promoting bf16 to f16. Quantized Florence-2 checkpoints are rejected at load with a named error rather than half-loaded. Neither half of this family has a quantized code path: the BART text stack and the DaViT tower are both built from dense `Linear` / `Embedding`, so a packed uint32 weight would reach MLX and abort the process instead of surfacing as an error. The check runs right after config parsing, before any weight file is read. ## Deviations and findings - **Upstream `image_feature_source` order is wrong for real checkpoints.** mlx-vlm's `ModelConfig` reads `image_feature_source`, `image_pos_embed`, and `visual_temporal_embedding` from the *top level* of `config.json`, but every real Florence-2 checkpoint stores them inside `vision_config`. Upstream therefore falls back to its dataclass defaults, and its `image_feature_source` default is `["temporal_avg_pool", "spatial_avg_pool"]`, the reverse of the checkpoint's `["spatial_avg_pool", "temporal_avg_pool"]`. That would put the pooled token at the end of the 577-token feature sequence instead of the front. mlxcel reads the config (as the HuggingFace Florence-2 implementation does) and refuses to guess when the field is absent; the parity run passes the checkpoint's own values to the reference so both sides execute the same recipe. A unit test pins the ordering behavior. - **The temporal embedding is not a no-op for still images.** With one frame it contributes row 0 of the cosine table, which is `sin(0), cos(0), ...` = `0, 1, 0, 1, ...`, not the zero vector. Skipping it would shift every odd channel by 1.0 before projection. Pinned by a unit test. - **`image_projection` is a bare tensor, not a `Linear`.** Upstream registers it as a raw `[in, out]` parameter applied as a right-hand matmul, so it is loaded and applied without a transpose or bias. - The issue text says `projection_dim: 768` is the fusion-stage width and the backbone emits `dim_embed[-1]`; the checkpoint config confirms both (1024 to 768). No correction needed there. - The joint attention mask reaches the encoder's self-attention only. Decoder cross-attention still passes no mask, matching the fact that this path drives a single image and a single prompt with no padding; a genuinely padded batch would need the mask threaded through cross-attention as well, which is not implemented here. The doc comment on `merge_input_ids_with_image_features` states this. - **The all-ones mask is skipped rather than applied.** `merge_input_ids_with_image_features` builds the mask the acceptance criteria ask for, but it is all ones by construction, and converting it would add a provably all-zero `[batch, 1, 1, seq]` tensor into every encoder layer's `[batch, heads, seq, seq]` logits (roughly 24M elementwise f16 adds per 585-token encode across 6 layers, all computing an identity). The general entry point is `encode_fused(inputs_embeds, attention_mask)`, which converts a caller-supplied mask; `encode_with_image_features` passes `None`. That removes provably redundant arithmetic rather than trading accuracy for speed, so no wall-clock figure is claimed and none was measured. `encoder_mask_is_identity_when_all_ones_and_load_bearing_otherwise` pins both directions: all-ones must equal `None`, and zeroing a position must change the output. - **Untrusted `config.json` integers are range-checked before they reach MLX.** MLX throws C++ exceptions that cannot cross the `cxx` bridge, so an out-of-range shape aborts the process instead of returning `Err`. `max_temporal_embeddings` is bounded (it drives an eager host allocation on the synthesis fallback), the loaded `embed_positions` table is checked to actually cover `max_position_embeddings` on both encoder and decoder, and `Florence2TextConfig` now rejects non-positive `d_model` / head counts / layer counts and out-of-vocabulary `pad_token_id` / `decoder_start_token_id`. `additive_attention_mask` returns `Result` instead of indexing an unchecked rank. - Runtime registration (`src/models/detection.rs`, `src/model_metadata.rs`, `docs/supported-models.md`, the `generate` / `generate_vlm` wiring, `mlxcel list`) is deliberately untouched: that is #856's scope. ## Test plan - [x] `cargo test --profile test-fast --features metal,accelerate --lib models::florence2::`: 51 passed. 37 in the new `florence2_fusion_tests.rs`, and 14 in `florence2_tests.rs` (7 pre-existing text-engine tests still green after the encoder mask change, plus 7 new config-validation tests). - [x] `cargo test --profile test-fast --features metal,accelerate --test florence2_fusion_parity`: 2 passed against the real `models/Florence-2-base-ft-bf16` checkpoint. Worst deviations vs the mlx-vlm reference: 2.2e-3 on image features (absmax 8.6), 1.5e-2 on encoder hidden states (absmax 45.8), 6.6e-3 on first-step logits (absmax 21.2). Mean and standard deviation agree to 1e-4 or better on all three, and the greedy token ids match exactly (`[0, 879, 27740, 868]`, decoding to `<s>unanswerable`, the sensible caption for a procedurally generated noise field). - [x] `cargo test --profile test-fast --features metal,accelerate --test florence2_text_parity --test florence2_vision_parity`: 3 passed, confirming the #852 and #853 pins are unmoved. - [x] `cargo clippy --profile test-fast --features metal,accelerate --lib --tests -- -D warnings`: clean. - [x] `cargo fmt --check`: clean. - [x] `python3 scripts/ci/check_cross_repo_refs.py`: OK. Closes #854

## Summary Lands the reusable BART-style seq2seq engine plus the Florence-2 text core (sub-issue of epic lablup#850): bidirectional encoder, causal decoder with encoder cross-attention, and the dual KV cache the decode loop needs. This is the foundation the vision-fusion (lablup#854), processor (lablup#855), and capstone (lablup#856) sub-issues drive the real model through. ## What changed - `src/models/florence2/layers.rs`: BART post-norm building blocks. `Florence2Attention` (bidirectional / causal / cross variants over one projection set), `Florence2EncoderLayer`, `Florence2DecoderLayer`, the `KvCache` pair, and an offset-aware additive causal mask. Post-norm ordering (`x = LN(x + Sublayer(x))`) per the mlx-vlm reference, unlike Whisper's pre-norm blocks. - `src/models/florence2/encoder.rs`: bidirectional encoder with learned absolute positions at the BART table offset of 2 and an embedding LayerNorm. Consumes pre-computed embeddings so Blaizzy/mlx-vlm-style fused image + text sequences enter the same path in lablup#854. - `src/models/florence2/decoder.rs`: causal decoder indexed by the running decode offset, cross-attending to the encoder output through per-layer dual caches. - `src/models/florence2/mod.rs`: `Florence2TextConfig` (parses the `text_config` sub-object of a real Florence-2 `config.json`, or a bare text config), `shift_tokens_right` with `decoder_start_token_id = 2`, `Florence2SeqCache`, and `Florence2TextModel` with the public engine API: `load` / `from_weights` (caller-supplied key prefix so the full-VLM loader reuses it), `embed_tokens`, `encode_tokens`, `encode_embeds`, `decode`, `make_cache`, `generate_greedy`. Weight loading filters to `language_model.*` tensors and applies the standard bf16 -> f16 conversion; a materialized `lm_head.weight` is preferred with a tied-shared-embedding fallback. - `src/models/florence2/florence2_tests.rs`: 13 unit tests on a tiny synthetic model covering config parsing, shift semantics, encoder bidirectionality, incremental-vs-full-sequence decode equality, prefill-plus-steps equality (offset causal mask), one-shot cross cache vs growing self cache, and greedy round-trip determinism. - `tests/florence2_text_parity.rs`: real-checkpoint parity against the mlx-vlm florence2 `LanguageModel` on `models/Florence-2-base-ft-bf16` (skips when the checkpoint is absent). Pins teacher-forced per-position argmax, logit slices at the first and last positions, and the greedy continuation. - `src/models/mod.rs`: module registration and re-exports; `src/lib/mlxcel-core/src/weights.rs`: `// Used by:` comment update for `load_weights_from_dir_filtered`. ## Correction to the issue text The issue body says 8 attention heads; the real checkpoint config has `encoder_attention_heads: 12` and `decoder_attention_heads: 12` (the mlx-vlm `TextConfig` dataclass default of 8 is overridden by the checkpoint). The implementation and tests derive every shape from the real `config.json` and pin 12 heads explicitly. ## Validation Reference values were captured by running the mlx-vlm florence2 `LanguageModel` (imported from the upstream source) on the same checkpoint with weights cast bf16 -> f16 to match mlxcel's Apple Silicon precision policy. mlxcel reproduces the reference exactly on teacher-forced argmax `[0, 2362, 16, 10, 2170, 13, 2]` and greedy continuation `[0, 2362]` (EOS stop), with logit slices within 0.05 in f16. ## Test plan - [x] `cargo test --profile test-fast --features metal,accelerate --lib florence2_tests::` (13 passed) - [x] `cargo test --profile test-fast --features metal,accelerate --test florence2_text_parity` (1 passed, real checkpoint) - [x] `cargo clippy --profile test-fast --features metal,accelerate --lib --tests -- -D warnings` (clean) - [x] `cargo fmt --check` (clean) - [x] `python3 scripts/ci/check_cross_repo_refs.py` (OK) Closes lablup#852

## Summary Implements the Florence-2 DaViT ("Dual Attention Vision Transformer") vision backbone, the pixel-side half of the model that will feed the BART seq2seq engine landed in lablup#1060. Four stages, each a strided `ConvEmbed` patch embedding followed by `depths[i]` blocks that alternate windowed spatial attention with channel attention, every half sandwiched between two depthwise positional convolutions. For the `Florence-2-base-ft` geometry and a 768x768 input the progression is 192x192 @ 128 -> 96x96 @ 256 -> 48x48 @ 512 -> 24x24 @ 1024, so the backbone emits `[B, 576, 1024]`. Scope is the backbone only. Fusion / projection, the processor, and end-to-end validation are separate sub-issues that depend on this one. ## What changed - `src/vision/encoders/florence2_davit.rs` (new): `Florence2VisionConfig` (parse + validate), the `sanitize` channels-last remap, `Florence2DaViT` with `from_weights` / `load` / `forward` / `forward_stages`, and `impl VisionEncoder`. - `src/vision/encoders/florence2_davit_blocks.rs` (new): `ConvEmbed`, `DepthWiseConv2d`, `Mlp`, `SpatialBlock`, `ChannelBlock`, `Block`. - `src/vision/encoders/florence2_davit_attention.rs` (new): `WindowAttention` and `ChannelAttention`. Split out of the blocks file to stay under the 500-line limit. - `src/vision/encoders/florence2_davit_tests.rs` (new): checkpoint-free unit tests for config parsing / validation, the sanitize pass and its idempotency, a synthetic-weight forward through the window pad / crop path, and the missing-weight error path. - `tests/florence2_vision_parity.rs` (new): numerical parity against the mlx-vlm reference on the real checkpoint, gated on checkpoint presence. - `src/vision/encoders/deepseekocr_sam.rs`: `window_partition` / `window_unpartition` reused rather than duplicated; their `// Used by:` comments now name Florence-2 DaViT, per `docs/code-guidelines.md`. - `src/vision/encoders/mod.rs`, `src/models/florence2/mod.rs`, `src/models/mod.rs`: module registration and re-export, so the fusion sub-issue can reach `Florence2DaViT` from the Florence-2 family directory alongside `Florence2TextModel`. ## API for the fusion sub-issue `Florence2VisionConfig::from_model_config(&Value)` (descends into `vision_config`, also accepts a bare vision-config object, mirroring `Florence2TextConfig::from_model_config`), `Florence2DaViT::from_weights(&WeightMap, &Florence2VisionConfig, prefix)` with `FLORENCE2_VISION_PREFIX = "vision_tower."`, `Florence2DaViT::load(&Path)` (config + `vision_tower.*` tensors + sanitize + bf16 -> f16 per the Apple Silicon precision policy), `forward` -> `[B, N, dim_embed[-1]]`, `forward_stages` -> per-stage `(tokens, (H, W))`, and `impl VisionEncoder`. ## Corrections to the issue text Every one of these was checked against `models/Florence-2-base-ft-bf16` (config.json and the actual safetensors keys and shapes) and against the upstream source, not inferred from the issue prose. 1. **`vision_config.model_type` is the empty string `""`, not `"davit"`.** The issue says to parse a `VisionConfig` with `model_type: davit`. Gating on that literal rejects every real Florence-2 checkpoint. Parsing here accepts `""`, `"davit"`, and an absent key (matching the upstream `VisionModel.__init__` guard, which allows `["davit", ""]`) and rejects anything else; the family is identified by the parent `model_type: florence2` plus the structural fields. Covered by a unit test that asserts all three spellings parse and a bogus one errors. 2. **There is no `hidden_size` in `vision_config`, and the backbone output is not 768.** The issue says "output image features with `hidden_size = 768`". The checkpoint has `projection_dim: 768` at both the top level and inside `vision_config`, while the last stage's `dim_embed` is `1024`, and `image_projection` is a `[1024, 768]` tensor. So the backbone emits `[B, N, 1024]`; the 1024 -> 768 projection is `image_projection` + `image_proj_norm` and belongs to the fusion sub-issue. `projection_dim` is parsed and retained on the config but never applied here. (Upstream's `VisionConfig` dataclass does carry a `hidden_size: int = 768` field, but `vision.py` never reads it.) 3. **The `window_partition` / `window_reverse` reuse pointer in the issue is wrong.** The issue says to reuse the patterns "from qwen2_5_vl and kimi_vl". Those encoders window differently. The helpers with matching semantics live only in `src/vision/encoders/deepseekocr_sam.rs`, and that is what this PR reuses (including the pad-inside-partition / crop-inside-unpartition contract, which lets the reference's explicit `mx.pad` and `x[:, :H, :W, :]` steps drop out). 4. **The issue omits four present, load-bearing config fields.** `num_groups: [4,8,16,32]`, `image_pos_embed: {type: "learned_abs_2d", max_pos_embeddings: 50}`, `visual_temporal_embedding: {type: "COSINE", max_temporal_embeddings: 100}`, and `image_feature_source: ["spatial_avg_pool","temporal_avg_pool"]`. Only `num_groups` affects this forward: it is `ChannelAttention.groups` per stage. The other three are consumed exclusively by `Florence2._encode_image` in the fusion stage, and upstream carries them on `ModelConfig` rather than `VisionConfig` while the real checkpoint puts them inside `vision_config`. All four are parsed and retained on `Florence2VisionConfig` so the fusion sub-issue reads them off one struct. 5. **The upstream `ChannelAttention` inline comment mis-states the attention shape.** It annotates `mx.matmul(q.transpose(0, 1, 3, 2), k)` as `(B, groups, N, N)`. The operands are `(B, g, C/g, N)` and `(B, g, N, C/g)`, so the product is `(B, groups, C/groups, C/groups)`: channel-by-channel, which is the whole point of channel attention. This port follows the code, and the doc comment on `ChannelAttention` records the discrepancy so nobody "corrects" it back. 6. **The `mlx-community` bf16 export already ships conv weights channels-last, so `sanitize` is a no-op on it.** `vision_tower.convs.0.proj.weight` is `[128, 7, 7, 3]` and `...dw.weight` is `[128, 3, 3, 1]` in the actual file, i.e. already `(O, kH, kW, I)`. The pass is still implemented in full (with the reference's `check_array_shape` guard on `convs.*.proj.weight` and the `shape[1] < shape[-1]` guard on `blocks.*.dw.weight`) so a raw PyTorch export converts correctly, and a unit test asserts both the transposing and the idempotent branches. 7. **`ChannelAttention` scales queries by the token count, not the head dimension.** `q = q * (float(N) ** -0.5)` where `N` is up to ~37k at stage 0. Unusual enough to look like a bug; it is not, and it is commented as deliberate. The softmax there runs in f32 and casts back, since the reduction axis is long. 8. **`DropPath` is training-only.** `drop_path_rate: 0.1` is present in the checkpoint and is parsed onto the config, but `self.training` is false at inference so the module is the identity. Implemented as a no-op with a comment saying why the config value does not reach the forward. `docs/supported-models.md` and `src/models/detection.rs` are deliberately untouched: this sub-issue does not make Florence-2 a loadable family (fusion does), so advertising it in the support table or wiring it into architecture detection now would claim a capability the binary does not have. Both belong with the fusion sub-issue. ## Validation **Input choice.** A deterministic synthetic pixel tensor, not a decoded image. Florence-2 image preprocessing is the processor sub-issue's scope, so reproducing it byte-identically in both runtimes is not possible yet, and what parity actually requires is that both sides see identical input bytes. Both sides build `x[0, c, h, w] = ((h * 768 + w) * 3 + c) % 251 / 251.0 - 0.5` in f32 (every step a single IEEE-754 operation over exactly represented integers, so the two implementations agree bit for bit) and then cast to f16. **Reference.** The mlx-vlm florence2 `VisionModel` (https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/models/florence2/vision.py) run in a scratch venv against `models/Florence-2-base-ft-bf16`, with `vision_tower.*` weights put through `VisionModel.sanitize` and cast bf16 -> f16 to match mlxcel's Apple Silicon precision policy. Captured per-stage shape, first-16 activations, and mean / std; the regeneration recipe is recorded in the test's module doc. **Tolerance.** Per-stage absolute, scaled to the activation magnitude at that stage, because both sides run f16 with different op ordering: `2e-3 / 1.2e-2 / 6e-2 / 1.2e-2` against stage peaks of `1.54 / 8.03 / 65.5 / 3.97`. Each is a small multiple of one f16 ulp at its stage's peak and at least 25x below that stage's own standard deviation. Nothing was loosened to make a failure pass; the tolerances were set after the first run already agreed. **Observed max absolute deviation** (Apple Silicon, `test-fast` profile): | Stage | Shape | Max abs deviation | Tolerance | Mean (ours / ref) | Std (ours / ref) | |-------|-------|-------------------|-----------|-------------------|------------------| | 0 | `[1, 36864, 128]` | 3.37e-4 | 2e-3 | -0.001657 / -0.001656 | 0.119353 / 0.119339 | | 1 | `[1, 9216, 256]` | 4.29e-3 | 1.2e-2 | 0.010804 / 0.010804 | 0.290855 / 0.290888 | | 2 | `[1, 2304, 512]` | 7.38e-4 | 6e-2 | 0.074498 / 0.074503 | 1.458913 / 1.458821 | | 3 | `[1, 576, 1024]` | 3.14e-3 | 1.2e-2 | -0.487351 / -0.487382 | 0.524581 / 0.524704 | ## Test plan - [x] `cargo test --profile test-fast --features metal,accelerate -p mlxcel --lib vision::encoders::florence2_davit` (7 passed) - [x] `cargo test --profile test-fast --features metal,accelerate -p mlxcel --test florence2_vision_parity -- --nocapture` (2 passed) - [x] `cargo test --profile test-fast --features metal,accelerate -p mlxcel --test florence2_text_parity` (1 passed, no regression from the re-export) - [x] `cargo clippy --profile test-fast --features metal,accelerate --lib --tests -- -D warnings` (clean) - [x] `cargo fmt --check` (clean) - [x] `python3 scripts/ci/check_cross_repo_refs.py` (clean) Closes lablup#853

…ng (lablup#1064) ## Summary Joins the DaViT vision tower (lablup#1063) to the BART seq2seq engine (lablup#1060) into `Florence2Model`, the assembled Florence-2 vision-language model, and completes weight loading for the whole checkpoint. Florence-2 concatenates rather than scattering: the tower output is pooled into a short feature sequence, projected into the text embedding space, and placed in front of the task-prompt embeddings to form the encoder input, with no image placeholder token in the prompt at all. ## What changed - `src/models/florence2/fusion.rs` (new): `LearnedPositionEmbedding2D` (learned column half first, then row half, over a square feature grid) and `PositionalEmbeddingCosine1D` (the `visual_temporal_embed.pos_idx_to_embed` table, taken from the checkpoint when present and synthesized from the closed form otherwise), plus `additive_attention_mask`, which turns a 0/1 mask into the `[batch, 1, 1, seq]` additive form attention consumes. - `src/models/florence2/checkpoint.rs` (new): `Florence2Config` (both halves plus `image_token_id`, defaulting to the text vocabulary size) and `sanitize`, which runs the DaViT conv channels-last remap over the whole weight map, drops the BART `final_logits_bias` buffer, and fills `model.shared.weight` from an `embed_tokens` copy for exports that materialize only the tied table. - `src/models/florence2/model.rs` (new): `Florence2Model` with `encode_image` / `encode_image_features` (`_encode_image`: 2-D grid position embedding, temporal embedding, `image_feature_source` pooling, `image_projection` 1024 to 768, `image_proj_norm`), `embed_prompt`, `merge_input_ids_with_image_features`, `encode` / `encode_with_image_features` / `encode_fused`, `make_cache`, `decode`, and `generate_greedy`. - `src/models/florence2/{encoder,layers}.rs`: the encoder and its blocks now take an optional additive attention mask, so the joint mask the fusion stage builds reaches the encoder's self-attention instead of being dropped. `encode_tokens` / `encode_embeds` pass `None`, leaving the text-only numerics unchanged; the new `Florence2TextModel::encode_embeds_with_mask` is what the fused path drives. - `src/models/florence2/florence2_fusion_tests.rs` (new): 37 unit tests covering both positional embeddings, the additive mask, all three `sanitize` passes, config parsing, load-time validation, and the whole `pixels -> features -> concat -> encoder -> decoder logits` chain on a tiny synthetic model (one block-free DaViT stage plus a one-layer BART stack), so the integrated path is exercised without a checkpoint. - `tests/florence2_fusion_parity.rs` (new): gated end-to-end parity against the mlx-vlm florence2 `Model` on `models/Florence-2-base-ft-bf16`, skipping cleanly when the checkpoint is absent. ## Precision and dtype handling The fused path pins its activation dtype from the text stack and casts the projected image features to it, and casts pixel input to the vision tower's own dtype, so a checkpoint whose two halves land on different dense dtypes does not silently promote the fused sequence to f32. Weight loading applies the usual Apple Silicon policy of promoting bf16 to f16. Quantized Florence-2 checkpoints are rejected at load with a named error rather than half-loaded. Neither half of this family has a quantized code path: the BART text stack and the DaViT tower are both built from dense `Linear` / `Embedding`, so a packed uint32 weight would reach MLX and abort the process instead of surfacing as an error. The check runs right after config parsing, before any weight file is read. ## Deviations and findings - **Upstream `image_feature_source` order is wrong for real checkpoints.** mlx-vlm's `ModelConfig` reads `image_feature_source`, `image_pos_embed`, and `visual_temporal_embedding` from the *top level* of `config.json`, but every real Florence-2 checkpoint stores them inside `vision_config`. Upstream therefore falls back to its dataclass defaults, and its `image_feature_source` default is `["temporal_avg_pool", "spatial_avg_pool"]`, the reverse of the checkpoint's `["spatial_avg_pool", "temporal_avg_pool"]`. That would put the pooled token at the end of the 577-token feature sequence instead of the front. mlxcel reads the config (as the HuggingFace Florence-2 implementation does) and refuses to guess when the field is absent; the parity run passes the checkpoint's own values to the reference so both sides execute the same recipe. A unit test pins the ordering behavior. - **The temporal embedding is not a no-op for still images.** With one frame it contributes row 0 of the cosine table, which is `sin(0), cos(0), ...` = `0, 1, 0, 1, ...`, not the zero vector. Skipping it would shift every odd channel by 1.0 before projection. Pinned by a unit test. - **`image_projection` is a bare tensor, not a `Linear`.** Upstream registers it as a raw `[in, out]` parameter applied as a right-hand matmul, so it is loaded and applied without a transpose or bias. - The issue text says `projection_dim: 768` is the fusion-stage width and the backbone emits `dim_embed[-1]`; the checkpoint config confirms both (1024 to 768). No correction needed there. - The joint attention mask reaches the encoder's self-attention only. Decoder cross-attention still passes no mask, matching the fact that this path drives a single image and a single prompt with no padding; a genuinely padded batch would need the mask threaded through cross-attention as well, which is not implemented here. The doc comment on `merge_input_ids_with_image_features` states this. - **The all-ones mask is skipped rather than applied.** `merge_input_ids_with_image_features` builds the mask the acceptance criteria ask for, but it is all ones by construction, and converting it would add a provably all-zero `[batch, 1, 1, seq]` tensor into every encoder layer's `[batch, heads, seq, seq]` logits (roughly 24M elementwise f16 adds per 585-token encode across 6 layers, all computing an identity). The general entry point is `encode_fused(inputs_embeds, attention_mask)`, which converts a caller-supplied mask; `encode_with_image_features` passes `None`. That removes provably redundant arithmetic rather than trading accuracy for speed, so no wall-clock figure is claimed and none was measured. `encoder_mask_is_identity_when_all_ones_and_load_bearing_otherwise` pins both directions: all-ones must equal `None`, and zeroing a position must change the output. - **Untrusted `config.json` integers are range-checked before they reach MLX.** MLX throws C++ exceptions that cannot cross the `cxx` bridge, so an out-of-range shape aborts the process instead of returning `Err`. `max_temporal_embeddings` is bounded (it drives an eager host allocation on the synthesis fallback), the loaded `embed_positions` table is checked to actually cover `max_position_embeddings` on both encoder and decoder, and `Florence2TextConfig` now rejects non-positive `d_model` / head counts / layer counts and out-of-vocabulary `pad_token_id` / `decoder_start_token_id`. `additive_attention_mask` returns `Result` instead of indexing an unchecked rank. - Runtime registration (`src/models/detection.rs`, `src/model_metadata.rs`, `docs/supported-models.md`, the `generate` / `generate_vlm` wiring, `mlxcel list`) is deliberately untouched: that is lablup#856's scope. ## Test plan - [x] `cargo test --profile test-fast --features metal,accelerate --lib models::florence2::`: 51 passed. 37 in the new `florence2_fusion_tests.rs`, and 14 in `florence2_tests.rs` (7 pre-existing text-engine tests still green after the encoder mask change, plus 7 new config-validation tests). - [x] `cargo test --profile test-fast --features metal,accelerate --test florence2_fusion_parity`: 2 passed against the real `models/Florence-2-base-ft-bf16` checkpoint. Worst deviations vs the mlx-vlm reference: 2.2e-3 on image features (absmax 8.6), 1.5e-2 on encoder hidden states (absmax 45.8), 6.6e-3 on first-step logits (absmax 21.2). Mean and standard deviation agree to 1e-4 or better on all three, and the greedy token ids match exactly (`[0, 879, 27740, 868]`, decoding to `<s>unanswerable`, the sensible caption for a procedurally generated noise field). - [x] `cargo test --profile test-fast --features metal,accelerate --test florence2_text_parity --test florence2_vision_parity`: 3 passed, confirming the lablup#852 and lablup#853 pins are unmoved. - [x] `cargo clippy --profile test-fast --features metal,accelerate --lib --tests -- -D warnings`: clean. - [x] `cargo fmt --check`: clean. - [x] `python3 scripts/ci/check_cross_repo_refs.py`: OK. Closes lablup#854

inureyes added type:enhancement New features, capabilities, or significant additions priority:medium Medium priority area:models Model architectures, weights, loading, metadata area:inference Generation, sampling, decoding (incl. speculative, DRY) status:review Under review labels Aug 7, 2026

inureyes added status:done Completed and removed status:review Under review labels Aug 7, 2026

inureyes merged commit baa1e13 into main Aug 7, 2026
8 checks passed

inureyes deleted the feature/issue-852-florence2-text-core branch August 7, 2026 03:38

inureyes mentioned this pull request Aug 7, 2026

feat(models): Florence-2 DaViT vision backbone #1063

Merged

6 tasks done

inureyes mentioned this pull request Aug 7, 2026

feat(models): Florence-2 vision-language fusion and full weight loading #1064

Merged

6 tasks done

This was referenced Aug 7, 2026

fix(loading): gemma4_unified quantized sanitize breaks dequantize/split equivalence #1074

Closed

epic: add Florence-2 (florence2) VLM support #850

Closed

inureyes self-assigned this Aug 27, 2026

Navigation Menu

Repository navigation

feat(models): Florence-2 BART seq2seq engine and text core - #1060

feat(models): Florence-2 BART seq2seq engine and text core#1060
inureyes merged 2 commits into
mainfrom
feature/issue-852-florence2-text-core

inureyes commented Aug 7, 2026

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

inureyes commented Aug 7, 2026

Summary

What changed

Correction to the issue text

Validation

Test plan

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant