Add NVIDIA LocateAnything-3B (MoonViT + Qwen2.5, autoregressive mode)

[beshkenadze]

Summary

Ports nvidia/LocateAnything-3B — a visual-grounding VLM (object detection / referring-expression grounding / pointing / GUI & text localization) — into mlx-vlm so it runs on Apple Silicon via mlx_vlm.generate.

Architecture: MoonViT-SO-400M vision tower (shared with Kimi-VL) + Qwen2.5-3B text backbone + a 2-layer MLP connector (mlp1). Output is structured coordinate tokens, e.g. <ref>remote</ref><box><64><152><273><244></box>, with coordinates quantized to <0>…<1000> (normalized).

What's implemented

• Model package mlx_vlm/models/locateanything/:
  • config.py — VisionConfig (moonvit) / TextConfig (qwen2) / ModelConfig with grounding token ids.
  • vision.py — MoonViT tower (2D RoPE, per-image block attention, 2×2 patch merge), ported from kimi_vl and reconciled to LocateAnything's weight names, PytorchGELUTanh activation, and LayerNorm eps.
  • language.py — standard Qwen2.5-3B causal LM (1D RoPE, GQA 16/2, tied embeddings) + the non-causal "magi" block mask for PBD.
  • locateanything.py — mlp1 connector + Model (vision → projector → scatter at <IMG_CONTEXT>) + sanitize().
  • pbd.py — Parallel Box Decoding (MTP block decoder + bbox decode utils).
  • image_processing_locateanything.py / processing_locateanything.py + chat_template.json (preprocessing matched bit-for-bit to the HF reference — bicubic ceil-resize, not center-crop).
• prompt_utils.py — register locateanything → LIST_WITH_IMAGE_FIRST.
• generate/dispatch.py — a small additive, opt-in hook routing fast/hybrid to PBD (gated on model_type == "locateanything"); slow and every other model are unchanged.
• tests/test_locateanything.py — config, vision/language/full-forward shapes, sanitize coverage, magi-mask, PBD decode utils, max-tokens, image-input handling (22 tests, all green).

Decoding modes

mode	description	throughput*
slow (default)	pure autoregressive	1×
fast	Parallel Box Decoding (MTP, parallel blocks)	~2×
hybrid	PBD with AR fallback on format irregularity	~2×

* 16-token COCO run. All three modes produce byte-identical grounding output.

Parity vs the PyTorch (CUDA) reference

Verified the MLX port numerically against the original HF/PyTorch model on an RTX 4090 (WSL, transformers==4.51, fp32). The HF vision_model + mlp1 were dumped on identical pixel_values, then fed to the MLX modules in fp32 and compared (scripts/la_parity_{ref,mlx}.py):

stage	grid 64×64 (no pos-emb interp)	grid 36×46 (with interp)
vision_model	cos 0.999937, mean|Δ| ≈ 1e-3	cos 0.9894
mlp1 (connector)	cos 0.999898	cos 0.9971

• The port's math is numerically faithful (cos ≈ 0.99994 on identical inputs).
• The entire residual on non-square grids comes from one op: the learnable 2D pos-emb bicubic interpolation — MLX's shared bicubic_interpolate kernel uses a = -0.5 vs PyTorch's a = -0.75. It is localized to the additive pos-emb (hence connector cos > vision cos: the connector's LayerNorm partly cancels it), affects every MLX MoonViT port (incl. kimi_vl), and does not change the grounding output. Tracked in Align bicubic_interpolate with PyTorch (a=-0.75, not -0.5) for MoonViT pos-emb parity #2.
• PBD parity: by the model's design invariant (hybrid falls back to AR for consistency), fast/hybrid are verified byte-identical to the slow (AR) path on the same input — so the verified AR path is itself the oracle for the parallel path.

Verification

• Unit tests: python -m unittest mlx_vlm.tests.test_locateanything → 22 passed.
• End-to-end (real weights, Apple Silicon):

  python -m mlx_vlm.generate --model nvidia/LocateAnything-3B \
    --image http://images.cocodataset.org/val2017/000000039769.jpg \
    --prompt "Detect all objects in the image." --max-tokens 128 --temperature 0.0
  

  → <ref>remote</ref><box><64><152><273><244></box><box><522><160><578><390></box> (boxes match the two remotes); prompt token count matches the HF processor exactly.
• Large images: fixed a vision-attention OOM (single-image dense [S,S] mask forced SDPA off the flash path; now flash, O(N) memory) — LocateAnything-3B: dense O(N²) vision attention OOMs on large images (Metal single-buffer cap) #3.
• Codex CLI review (codex exec review): no P0/P1; two P2 edge-cases (PBD max_tokens, mx.array image input) fixed.

Blast radius

Everything PBD/model-specific lives in the locateanything package; the only shared edits are one additive line in prompt_utils.py and one additive, gated hook in dispatch.py. No other model, the vision tower, the processor, or shared SDPA is affected; slow AR remains the default.

Quantized weights

MLX builds published to mlx-community: bf16, 8bit, 4bit (mixed 4/8-bit — pure 4-bit degrades the tied coordinate-token embedding).

[beshkenadze]

Superseded by the upstream PR → Blaizzy#1242 (same branch). Continuing review there.