feat: Derive worst-case image context floors for the remaining qualified OpenXLA vision families

[inureyes]

Problem / Background

xla_image_context_floor in src/multimodal/host_preprocessor.rs derives, from config.json and preprocessor_config.json alone with no weights loaded, the largest number of logical prompt tokens that one image can expand into on a given checkpoint. ensure_xla_image_context_capacity in the same file uses that number to reject at startup a static OpenXLA graph shape that could never admit an image, instead of letting every image request run its whole vision tower and only then fail at admission.

Only Molmo2 has a formula today. xla_image_context_floor matches ModelType::Molmo2VLM and returns None for everything else. LlavaVLM and Qwen2VL are both qualified for the OpenXLA image path (load_xla_image_preprocessor in the same file constructs LlavaIreeHostPreprocessor and Qwen2VlIreeHostPreprocessor for them), so they fall through the guard silently and still hit the late-admission failure mode the guard exists to prevent.

Current Behavior

xla_image_context_floor dispatches on crate::models::get_model_type and delegates to molmo2_image_context_floor for Molmo2VLM; every other family returns None. ensure_xla_image_context_capacity treats None as "not derived for this checkpoint", not "no images", and returns Ok(()), which is the right conservative behavior for a family without a formula but leaves LLaVA and Qwen2-VL unguarded.

Proposed Solution

Add a per-family derivation for LLaVA and Qwen2-VL, following the Molmo2 one as the worked example.

The Molmo2 formula, as the pattern to follow. The prompt carries one pooled token per pooled patch, plus one column token per pooled row, for a low-resolution crop and a high-resolution tiling:

tokens = lo_h * (lo_w + 1) + hi_h * (hi_w + 1) + 4

The low-resolution crop is always one tile. The high-resolution tiling is whichever (rows, cols) the processor picks for the image, subject to rows * cols <= max_crops. The worst case is found by walking every admissible tiling rather than assuming the squarest or the largest-area one, because the per-row column token makes tall tilings more expensive than wide ones of equal area. The four trailing tokens are the low-resolution start and end and the high-resolution start and end. On the pinned 4B checkpoint at /home/inureyes/models/mlx/molmo2-4b (max_crops = 8, 378px crops, 14px patches, 2x2 pooling) that gives 424 tokens for a square image and 1834 for a tall one, which is why sizing on the observed square case would still reject ordinary photographs.

LLaVA. Base LLaVA contributes a fixed number of tokens per image, derived from the vision tower's image size, patch size, and whether the CLS token is kept. The llava_next style multi-patch variants add an anyres grid on top, so the floor has to walk the configured grid pinpoints the way the Molmo2 derivation walks tilings, plus whatever separator or newline tokens the variant inserts per row. The derivation must read the actual variant from config rather than assuming the base case.

Qwen2-VL. The bound is set by the smart-resize policy plus the spatial merge: the number of visual tokens is the resized pixel area divided by patch_size squared and then by merge_size squared, so the worst case is the maximum admitted pixel count under the resize policy, rounded to the patch and merge granularity. The floor therefore depends on max_pixels (or its equivalent in preprocessor_config.json) and on how smart-resize rounds, not on any particular input image.

Scope

In scope: src/multimodal/host_preprocessor.rs (new llava_image_context_floor and qwen2_vl_image_context_floor alongside molmo2_image_context_floor, plus the new arms in xla_image_context_floor), and src/multimodal/host_preprocessor_tests.rs for the pinned tests.

Out of scope: families that are not qualified for the OpenXLA image path (Muse Glimmer VLM is explicitly rejected in load_xla_image_preprocessor), any change to ensure_xla_image_context_capacity's policy, and bucketing the capacity so that a large floor is not a throughput tax (that is #1271).

Implementation Notes

• Reuse: follow molmo2_image_context_floor exactly in shape. It reads the JSON with serde_json, returns None on any missing or unreadable key rather than guessing, rejects zero denominators before dividing, and uses checked arithmetic throughout so a malformed config cannot panic or overflow into a bogus floor. Do not introduce a different error convention per family.
• Constraints: the derivation runs on the server path before any preprocessor exists, so it must not load weights, construct a processor, or touch the vision tower. Config files only.
• Edge cases: a missing preprocessor_config.json must report None (there is already a test for this on the Molmo2 path), a config whose variant is not recognized must report None rather than falling back to the base-case formula, and a config with a zero or absent patch or merge size must report None instead of dividing.
• Error handling: None is the only failure signal. It means the guard stays silent for that checkpoint, which is the existing conservative default. Nothing here should return an error or log a warning on a config it cannot read.
• Overestimating is safe, underestimating is not. A floor that is too high costs the operator a larger graph. A floor that is too low reinstates exactly the late-admission failure the guard exists to prevent, so where the worst case is ambiguous, round up.

Acceptance Criteria

• xla_image_context_floor returns Some(floor) for a LlavaVLM checkpoint and for a Qwen2VL checkpoint, derived from config alone with no weights loaded.
• A unit test per family pins the number for a real checkpoint's geometry, in the style of molmo2_image_floor_is_the_tallest_tiling_not_the_square_one in src/multimodal/host_preprocessor_tests.rs, and asserts the worst case rather than a convenient case.
• For each family, the derived floor is validated against a real checkpoint by preprocessing an actual image of the worst-case shape and confirming it produces no more logical prompt tokens than the floor.
• A checkpoint with a missing or unreadable preprocessor_config.json still returns None for both new families, with a test.
• ensure_xla_image_context_capacity rejects at startup, for both new families, a capacity one token below the derived floor and accepts it at the floor, mirroring the existing Molmo2 boundary test.

Verification

cargo test --release --features cuda,xla-iree \
  -p mlxcel host_preprocessor -- --test-threads=1

eval "$(bash scripts/iree/setup-cuda.sh --env)"
MLX_ENABLE_TF32=0 MLXCEL_BACKEND=xla MLXCEL_XLA_DEVICE=cuda \
  ./target/release/mlxcel generate -m <llava-or-qwen2vl checkpoint> --image <worst-case-shape image> -p "Describe this image." -n 16

A pass is: the pinned unit tests green, and the real-checkpoint run reporting a prompt token count at or below the derived floor for a worst-case-shaped image. CUDA test runs must use --test-threads=1.

Technical Considerations

The floor's value grows once #1271 lands: with a single static graph shape a large floor is a throughput tax on every text-only request, so the guard's message today has to offer the operator a bad choice. With bucketing, the floor becomes the size of the image bucket and costs text requests nothing, which makes deriving it correctly for every qualified family more useful, not less.

Related: PR #916, #871, #1271.