feat: add shared DINOv2 feature encoder

TECHNICAL_REPORTS/2549-dinov2-backbone-20261023.en.md

@@ -0,0 +1,127 @@
+# PR #2549: Add Shared DINOv2 Feature Encoder
+
+**Date**: 2026-10-23
+
+**Status**: Local validation complete (PR open; final CI gate pending)
+
+**Languages**: Rust, Python, Markdown
+
+**Risk**: Medium
+
+## Executive Summary
+
+PR #2549 adds a reusable DINOv2 vision backbone for standard and register-token Hugging Face checkpoints and exposes standalone image feature extraction through `mlxcel predict`. It implements checkpoint conversion, positional interpolation, mask handling, normalized intermediate features, the giant preset's SwiGLU path, checkpoint-driven preprocessing, and three output tensors: `last_hidden_state`, `pooler_output`, and `hidden_patch_tokens`.
+
+The final source passed formatting-aware clippy, eight repository contract targets, 13 focused tests with genuine CPU f32 checkpoint references, and guarded ROCm CLI runs for both supported checkpoint layouts on gfx1151. Independent implementation and security/performance reviews have no remaining findings. GitHub CI had no reported failures at this report cut, but three central compile/lint jobs were still queued, so merge readiness remains subject to that gate.
+
+## 1. Problem Statement
+
+Video Depth Anything, MoGe-3, and SAM 3D Objects depend on a DINOv2 transformer, but mlxcel had no reusable implementation. Porting a private transformer into each dense model would duplicate checkpoint conversion, positional interpolation, token handling, and numerical validation. Public DINOv2 checkpoints also needed a standalone feature path that does not pass through text generation or object detection.
+
+The implementation had to support both public checkpoint layouts, exclude register tokens from spatial output, reproduce both upstream positional interpolation modes, and preserve the Hugging Face image processor contract. Incorrect resize rounding, convolution layout conversion, mask dtype, or token slicing changes every downstream feature and can silently invalidate dependent model ports.
+
+## 2. Change Summary
+
+| Area | Change |
+|---|---|
+| Shared encoder | Added DINOv2 patch embedding, class/register/mask tokens, attention, LayerScale blocks, dense MLP and giant SwiGLU FFN, final normalization, and normalized intermediate-layer extraction. |
+| Checkpoint loading | Added standard/register Hugging Face key remapping, fused QKV and SwiGLU projections, guarded OIHW-to-OHWI patch-convolution conversion, scoped parent-checkpoint loading, and shape validation. |
+| Position handling | Added rectangular bicubic position interpolation in Hugging Face size mode with configurable antialiasing and legacy scale-factor mode with the `0.1` offset. |
+| Standalone prediction | Registered `dinov2` and `dinov2_with_registers` for the `features` task and writes normalized full sequence, class pooling, and spatial-only patch tensors. |
+| Preprocessing | Reads `preprocessor_config.json` and reproduces the slow Hugging Face processor: RGB conversion, shortest-edge resize, Pillow-compatible bicubic uint8 rounding, center crop, rescaling, and channel normalization. |
+| Safety | Bounds architecture dimensions, processor sizes, source and resized pixels, batch token counts, register tokens, and intermediate-layer requests before expensive graph operations. Public attention input and public attention/MLP construction are validated. |
+| Fixtures and docs | Added pinned full-output fixtures for both small checkpoint variants, a reproducible reference generator, environment-variable documentation, architecture capability metadata, and dense-prediction operator guidance. |
+
+The implementation diff before this report is 19 files with 2,829 insertions and 14 deletions. Its commits are `e3ab2c2e`, `29b30a27`, `0b0ea0c3`, and `bb89237b`.
+
+## 3. Technical Decisions
+
+### 3.1 One reusable backbone for dense consumers
+
+`Dinov2` owns the checkpoint-compatible transformer, while `Dinov2Encoder` adds resize, ImageNet normalization, selected intermediate layers, and spatial grid reshaping for dense models. Public `Attention` and `Mlp` types let SAM 3D Objects reuse the same verified primitives. Scoped construction copies only a named DINOv2 subtree, leaving the parent weight map available for sibling heads.
+
+### 3.2 Preserve upstream token and interpolation contracts
+
+The sequence order is class token, optional register tokens, then patch tokens. Position embeddings apply to class and patches before registers are inserted, matching Hugging Face. `forward_features` retains all tokens, `pooler_output` uses the normalized class token, and spatial outputs exclude both class and registers. Position interpolation supports the explicit-size antialias mode used by Hugging Face checkpoints and the legacy scale-factor offset mode used by dense consumers.
+
+### 3.3 Match the slow Hugging Face processor at the byte boundary
+
+The standalone predictor uses checkpoint processor settings instead of assuming patch-aligned input. Its bicubic path reproduces Pillow's fixed-point coefficient and uint8 rounding behavior before rescaling and normalization. This choice was verified against output from the actual slow processor.
+
+### 3.4 Validate checkpoint dimensions before allocation
+
+Configuration and public layer APIs reject invalid or excessive hidden sizes, layer counts, patch grids, register counts, image sizes, and intermediate requests. Attention uses the shared safe dispatch; `Attention::forward` rejects invalid input shapes, while the public attention and MLP constructors validate dimensions and checkpoint weights. These checks keep malformed checkpoint metadata from creating oversized host buffers or MLX graphs.
+
+### 3.5 Commit complete reference tensors
+
+The two safetensors fixtures contain normalized NHWC input and all three expected output tensors rather than sparse probes. Each fixture records its public model ID, immutable checkpoint revision, library versions, dtype, device, and input convention. Future consumers therefore have a stable end-to-end CPU oracle at the issue's relative RMS tolerance of `1e-3`.
+
+## 4. Architecture and Data Flow
+
+```text
+source image
+    -> checkpoint BitImageProcessor policy
+    -> f32 NHWC pixels
+    -> patch convolution + class/register/mask token handling
+    -> DINOv2 blocks with LayerScale
+    -> final LayerNorm
+    -> last_hidden_state
+       + normalized class pooler_output
+       + register-free hidden_patch_tokens
+```
+
+Dense consumers can instead provide `[0,1]` NHWC images and a requested token grid to `Dinov2Encoder`. That path performs antialiased bilinear resizing, ImageNet normalization, selected normalized intermediate layers, and reshaping of patch tokens into NHWC feature grids.
+
+## 5. Verification
+
+Final source head `bb89237b` passed `cargo clippy --profile test-fast --lib --tests -- -D warnings` with status 0 in 2 minutes 29 seconds. All eight fast contract targets and their self-tests also passed with status 0. Independent implementation and security/performance reviews found no remaining issues after fixes for exact processor behavior, mask dtype, scoped prefix APIs, allocation bounds, public attention validation, shared fast dispatch, bounded intermediates, and fixture integrity.
+
+The final focused DINOv2 suite rebuilt in 8 minutes 39 seconds, then completed 13 tests in 49.71 seconds with 13 passed, 0 failed, 0 ignored, and 9,888 filtered. Both full checkpoint comparisons ran on CPU f32 rather than skipping.
+
+| CPU f32 reference | `last_hidden_state` RMS | `pooler_output` RMS | `hidden_patch_tokens` RMS |
+|---|---:|---:|---:|
+| `facebook/dinov2-small` | `2.432286715e-6` | `1.959671192e-6` | `2.433967071e-6` |
+| `facebook/dinov2-with-registers-small` | `4.238213112e-6` | `2.756216562e-6` | `4.275870861e-6` |
+
+The slow-processor input comparison reached relative RMS `1.197926287e-7`. The standalone standard-checkpoint predictor reached `2.705206368e-6`, `2.218856821e-6`, and `2.706949506e-6` for the full sequence, pooled class token, and spatial patch tokens. These measurements also ran on CPU f32; the suite's serialized ROCm guard only coordinates shared hardware access.
+
+The CLI built with status 0 in 8 minutes 10 seconds. A single guarded run with `MLXCEL_DEVICE=gpu` bound the gfx1151 AMD Radeon 8060S and exercised both public small checkpoint variants with status 0.
+
+| ROCm CLI checkpoint | Output and shape | Relative RMS |
+|---|---|---:|
+| Standard | `last_hidden_state` f32 `[1,257,384]` | `2.844445329e-6` |
+| Standard | `pooler_output` f32 `[1,384]` | `2.473357604e-6` |
+| Standard | `hidden_patch_tokens` f32 `[1,256,384]` | `2.845811506e-6` |
+| Register | `last_hidden_state` f32 `[1,261,384]` | `4.548628530e-6` |
+| Register | `pooler_output` f32 `[1,384]` | `2.131596900e-6` |
+| Register | `hidden_patch_tokens` f32 `[1,256,384]` | `4.596212478e-6` |
+
+Every ROCm result is below `1e-3`. Each variant wrote `meta.json` and `result.safetensors`; the safetensors contain exactly the three documented keys with the shapes and f32 dtype shown above. `arch --json` for both model types and the `predict`, image, and embeddings paths also passed.
+
+At the report cut, GitHub CI showed no failures. The central `cargo-clippy`, CPU-only link, and OpenXLA feature compile jobs were still queued, so CI is not recorded as complete and remains the merge gate.
+
+## 6. Limits
+
+The full public giant checkpoint was not executed. Giant preset values and SwiGLU remapping/sanitizer behavior have unit coverage, but this report makes no genuine giant forward, real-weight parity, or memory claim. Metal runtime parity, DINO training/classification heads, `/v1/embeddings`, and the RF-DETR windowed variant remain outside this PR. No throughput benchmark was performed.
+
+## 7. Learning Points
+
+- A shared backbone should expose both full-sequence and spatial-only views because register tokens are valid transformer context but are not image-grid locations.
+- Image processor parity includes integer resize behavior before normalization; matching only the floating-point network does not establish end-to-end parity.
+- Scoped checkpoint loading lets composite dense models reuse a backbone without consuming or renaming unrelated head weights.
+- Full-tensor fixtures reveal systematic drift that sparse probes can miss and preserve enough metadata to reproduce the oracle later.
+
+## 8. Follow-up Actions
+
+- Let the central CI gate complete before merge.
+- Exercise a real giant checkpoint when suitable hardware and memory are available.
+- Validate Metal runtime parity on suitable Apple hardware.
+- Reuse this encoder in the dependent Video Depth Anything, MoGe-3, and SAM 3D Objects ports.
+
+## References
+
+- PR #2549: `feat: add shared DINOv2 feature encoder`
+- Issue #2350: DINOv2 backbone for dense models
+- Epic #2382: dense-prediction model delivery
+- `docs/dense-prediction.md`: standalone feature output and processor contract
+- `tests/fixtures/dinov2/README.md`: fixture provenance and regeneration