feat(xla): support Molmo2 indexed attention pooling #871

Parent and dependency

• Parent epic: epic: OpenXLA multimodal / VLM architecture support (Qwen2-VL, Gemma3n, Phi4MM, Molmo, and more) #566
• Blocked by test(xla): validate a LLaVA reference architecture end to end #862
• Existing MLX implementation: src/vision/molmo2_vl.rs, src/vision/encoders/molmo2.rs, src/vision/processors/molmo2.rs, and src/models/molmo2.rs

Problem

Molmo2 has its own visual/data contract: pre-flattened patches use a Linear patch embedding, position embeddings depend on the patch grid, configured intermediate ViT layers are concatenated, image_token_pooling contains patch indices with negative invalid sentinels, attention pooling produces one vector per output token, and a SwiGLU projector is additively merged at scanned image_patch_id positions.

This differs from Molmo v1's processor-supplied image_input_idx mapping and must not share the wrong merge/index policy.

Goal

Support a pinned Molmo2 checkpoint end to end on XLA with its exact patch-grid, layer-selection, indexed attention-pooling, additive merge, and continuous-batch behavior.

Required implementation

• Reuse/share Molmo2Processor outputs: flattened crops/patches, image_token_pooling, image_grids, and image_num_crops. Preserve integer dtypes, shapes, and negative pooling sentinels in the prepared payload.
• Emit Linear patch embedding, configured position embedding/interpolation policy, every required ViT block, exact selected-layer concatenation, indexed 2D attention pooling, optional pooling attention mask, and SwiGLU image projector through StableHLO/IREE.
• Implement batched gather for pooling indices with invalid entries clamped only for safe gather and then masked out. The query mean denominator must count valid patches and must not divide by zero for an all-invalid group.
• Validate crop/patch/grid counts, position-table bounds/interpolation support, layer indices, pooling index ranges, pool-group sizes, projected-token count, and image_patch_id count before native invocation.
• Scan logical expanded tokens for the configured image_patch_id, preserve position order, and add one projected feature per position. Do not replace text embeddings and do not use Molmo v1 image_input_idx semantics.
• Include position policy, selected layers, pooling-mask mode, projector config, token IDs, and static bucket sizes in artifact compatibility identity.
• Integrate filtered loading, capability detection, CLI/server batching, cancellation, slot reset, and no-media text fallback without a duplicate full MLX decoder.

Non-goals

• Molmo v1 sparse index mapping.
• Unsupported position interpolation beyond the pinned checkpoint unless independently validated.
• Molmo Point task-specific coordinate decoding.

Validation

• Pin one Molmo2 checkpoint/revision and compare patches/grids, position embeddings, selected ViT states, concatenated features, safe gather plus valid mask, pooling queries/outputs, projected features, scanned positions, merged embeddings, logits/KV, and greedy tokens against MLX and an independent oracle.
• Cover multiple crops, partially and fully invalid pooling groups, exact/default and nondefault supported grids, wrong grid/crop metadata, feature/token-count mismatch, no-media text, cancellation, and slot reuse.
• Add negative fixtures for unmasked clamped indices, wrong denominator, replacement merge, and wrong selected layers.
• Run CLI and mixed continuous-batch server tests on a production-relevant IREE target, plus test(xla): validate a LLaVA reference architecture end to end #862 and text-only Molmo2 gates.

Acceptance criteria

• Molmo2's Linear patch/position/layer-selection path executes through IREE.
• Negative pooling entries cannot leak clamped patch zero into pooling results.
• Pooling and additive image_patch_id merge match both reference implementations.
• Deterministic output is token-exact on the pinned fixture.
• Invalid/all-invalid groups, grid mismatch, and count mismatch fail safely.
• Molmo v1 semantics are not reused where contracts differ.
• Text-only Molmo2 behavior remains token-exact.