perf(vision/mllama): build cross-attention states from real tiles only

src/models/mllama/text.rs

@@ -581,4 +581,113 @@ mod tests {
             "a new image must produce a different cached key, not a stale one"
         );
     }
+
+    // --- Real-tile cross-states equivalence (issue #527 perf follow-up). ---
+    //
+    // The reference masks every padding-tile position out of the text
+    // cross-attention with an additive -1e9 (`cross_attention_mask` derived
+    // from the per-image `num_tiles`). `exp(logit - 1e9)` underflows to
+    // exactly 0.0, so those positions add exact zeros to the softmax numerator
+    // and denominator: attending over the REAL-tile rows alone is the same
+    // computation. These tests pin that equivalence at the byte level, which
+    // is what licenses `MllamaVLModel` to drop the padding-tile rows from
+    // `cross_attention_states` instead of threading a reference-style mask.
+
+    /// `[1, rows, HIDDEN]` cross-states with identifiable per-row content.
+    fn cross_rows(rows: i32, seed: usize) -> UniquePtr<MlxArray> {
+        mlxcel_core::from_slice_f32(&fill((rows * HIDDEN) as usize, seed), &[1, rows, HIDDEN])
+    }
+
+    /// Additive `[1, 1, Q_LEN, kv]` mask: 0 everywhere, -1e9 on `masked` keys
+    /// (the reference's padding-tile columns).
+    fn additive_kv_mask(kv: i32, masked: &[i32]) -> UniquePtr<MlxArray> {
+        let mut vals = vec![0.0f32; (Q_LEN * kv) as usize];
+        for q in 0..Q_LEN {
+            for &k in masked {
+                vals[(q * kv + k) as usize] = -1e9;
+            }
+        }
+        mlxcel_core::from_slice_f32(&vals, &[1, 1, Q_LEN, kv])
+    }
+
+    fn rows_slice(x: &MlxArray, start: i32, end: i32) -> UniquePtr<MlxArray> {
+        mlxcel_core::slice(x, &[0, start, 0], &[1, end, HIDDEN])
+    }
+
+    /// Single image, 1 real tile of 4 (2 patch rows per tile): attention over
+    /// the real rows only is byte-identical to reference-masked attention over
+    /// all rows, garbage padding-lane features included.
+    #[test]
+    fn real_tile_rows_match_reference_masked_full_rows() {
+        let config = tiny_config();
+        let layer = build_layer(&config);
+        let h = hidden_states();
+
+        // 4 tiles x 2 patches = 8 kv rows; rows 2..8 are padding-lane features
+        // (arbitrary non-zero values, as the tower genuinely produces).
+        let full = cross_rows(8, 42);
+        let mask = additive_kv_mask(8, &[2, 3, 4, 5, 6, 7]);
+        let masked_out = layer.forward(&h, &full, Some(&mask));
+
+        // The fill-once KV cache belongs to the previous cross-states.
+        layer.invalidate_kv_cache();
+        let real_only = rows_slice(&full, 0, 2);
+        let sliced_out = layer.forward(&h, &real_only, None);
+
+        assert_eq!(
+            max_abs_diff(&masked_out, &sliced_out),
+            0.0,
+            "dropping the -1e9-masked padding-tile rows must be byte-identical"
+        );
+    }
+
+    /// Ragged multi-image (media 0: 1 real tile of 2, media 1: 2 of 2):
+    /// media-major concatenation of each image's real rows is byte-identical
+    /// to reference-masked attention over the full row set.
+    #[test]
+    fn ragged_real_tile_rows_match_reference_masked_full_rows() {
+        let config = tiny_config();
+        let layer = build_layer(&config);
+        let h = hidden_states();
+
+        // [media0 t0, media0 t1(pad), media1 t0, media1 t1] x 2 patches.
+        let full = cross_rows(8, 77);
+        let mask = additive_kv_mask(8, &[2, 3]);
+        let masked_out = layer.forward(&h, &full, Some(&mask));
+
+        layer.invalidate_kv_cache();
+        let media0 = rows_slice(&full, 0, 2);
+        let media1 = rows_slice(&full, 4, 8);
+        let sliced = mlxcel_core::concatenate(&media0, &media1, 1);
+        let sliced_out = layer.forward(&h, &sliced, None);
+
+        assert_eq!(
+            max_abs_diff(&masked_out, &sliced_out),
+            0.0,
+            "ragged per-image selection must be byte-identical to the \
+             reference-masked full attention"
+        );
+    }
+
+    /// All tiles real: the all-zero reference mask is the no-mask computation;
+    /// there is nothing to drop and no behavior change.
+    #[test]
+    fn full_tile_mask_is_the_unmasked_computation() {
+        let config = tiny_config();
+        let layer = build_layer(&config);
+        let h = hidden_states();
+
+        let full = cross_rows(8, 99);
+        let mask = additive_kv_mask(8, &[]);
+        let masked_out = layer.forward(&h, &full, Some(&mask));
+
+        layer.invalidate_kv_cache();
+        let unmasked_out = layer.forward(&h, &full, None);
+
+        assert_eq!(
+            max_abs_diff(&masked_out, &unmasked_out),
+            0.0,
+            "an all-zero mask must not perturb the attention output"
+        );
+    }
 }

src/vision/encoders/mllama.rs

@@ -670,6 +670,7 @@ mod tests {
     use super::{MllamaVisionModel, prepare_aspect_ratio_attention_mask};
     use crate::models::mllama::MllamaVisionConfig;
     use mlxcel_core::weights::WeightMap;
+    use mlxcel_core::{MlxArray, UniquePtr};
 
     fn value_at(mask: &mlxcel_core::MlxArray, i: i32, j: i32) -> f32 {
         // mask: [1, 1, L, L]
@@ -893,4 +894,237 @@ mod tests {
             "a dense tower must not report quantized projections"
         );
     }
+
+    // --- Padding-tile semantics (issue #527 perf follow-up). ---
+    //
+    // These tests document WHY the tower must keep processing the zero-padding
+    // tiles even though they look like wasted compute. The aspect-ratio mask is
+    // an outer product of the inverted validity vector, so it adds `-1e9` only
+    // where BOTH the query and the key are padding positions. Real-tile queries
+    // therefore attend to padding-tile keys with a zero bias (strictly positive
+    // softmax weight), and padding-tile lanes ingest real-tile content in one
+    // layer and feed it back to real-tile queries in the next. The padding
+    // tiles act as trained register lanes: dropping them from the tower input
+    // changes the real tiles' output. This matches the reference exactly
+    // (mlx-vlm `_prepare_aspect_ratio_attention_mask` and the HF processor,
+    // which always pads the tile axis to `max_image_tiles` before the tower).
+
+    /// Real->padding and padding->real attention is UNMASKED; only
+    /// padding->padding pairs carry the -1e9 bias. This is the structural
+    /// reason a "run the tower on real tiles only" optimization is unsound.
+    #[test]
+    fn mask_keeps_real_to_padding_attention_open() {
+        // 2 tiles, 2 patches each; tile 1 is a padding tile.
+        // Flattened positions: 0,1 = tile 0 (real), 2,3 = tile 1 (padding).
+        let ar_mask = mlxcel_core::from_slice_i32(&[1, 0], &[1, 2]);
+        let mask = prepare_aspect_ratio_attention_mask(&ar_mask, 2, 2, 2);
+
+        // Real query -> padding key: open (bias 0), NOT excluded.
+        assert_eq!(value_at(&mask, 0, 2), 0.0);
+        assert_eq!(value_at(&mask, 1, 3), 0.0);
+        // Padding query -> real key: also open.
+        assert_eq!(value_at(&mask, 2, 0), 0.0);
+        assert_eq!(value_at(&mask, 3, 1), 0.0);
+        // Only padding query -> padding key is masked.
+        assert!(value_at(&mask, 2, 3) < -1e8);
+        assert!(value_at(&mask, 3, 2) < -1e8);
+    }
+
+    // --- Forward-capable tiny tower harness. ---
+
+    /// Tiny but forward-runnable tower: 4x4 image, patch 2 (4 patches + cls =
+    /// 5, padded to 8), hidden 8, 2 heads, 2 local + 1 gated global layer,
+    /// max_num_tiles 4.
+    fn forward_vision_config() -> MllamaVisionConfig {
+        serde_json::from_str(
+            r#"{
+                "image_size": 4,
+                "patch_size": 2,
+                "hidden_size": 8,
+                "intermediate_size": 16,
+                "num_hidden_layers": 2,
+                "num_global_layers": 1,
+                "num_attention_heads": 2,
+                "max_num_tiles": 4,
+                "intermediate_layers_indices": [0]
+            }"#,
+        )
+        .expect("forward-capable tiny mllama vision config")
+    }
+
+    /// Deterministic pseudo-random fill in roughly `[-0.5, 0.5]`.
+    fn fill(n: usize, seed: usize) -> Vec<f32> {
+        (0..n)
+            .map(|i| ((i * 131 + seed * 977 + 7) % 251) as f32 / 251.0 - 0.5)
+            .collect()
+    }
+
+    fn put(map: &mut WeightMap, key: &str, shape: &[i32], seed: usize) {
+        let n: i32 = shape.iter().product();
+        map.insert(
+            key.to_string(),
+            mlxcel_core::from_slice_f32(&fill(n as usize, seed), shape),
+        );
+    }
+
+    /// Full dense weight set with non-degenerate values so attention mixing is
+    /// observable (the zero-filled loader-test weights would hide it).
+    fn build_forward_weights(prefix: &str) -> WeightMap {
+        let (h, inter, np, tiles, ar_rows) = (8, 16, 5, 4, 9);
+        let mut w = WeightMap::new();
+
+        put(&mut w, &format!("{prefix}.class_embedding"), &[h], 1);
+        // PyTorch conv layout [out, in, kH, kW]; the loader transposes it.
+        put(
+            &mut w,
+            &format!("{prefix}.patch_embedding.weight"),
+            &[h, 3, 2, 2],
+            2,
+        );
+        for (name, seed) in [("layernorm_pre", 3), ("layernorm_post", 5)] {
+            put(&mut w, &format!("{prefix}.{name}.weight"), &[h], seed);
+            put(&mut w, &format!("{prefix}.{name}.bias"), &[h], seed + 1);
+        }
+
+        put(
+            &mut w,
+            &format!("{prefix}.gated_positional_embedding.embedding"),
+            &[np, h],
+            7,
+        );
+        put(
+            &mut w,
+            &format!("{prefix}.gated_positional_embedding.gate"),
+            &[1],
+            8,
+        );
+        put(
+            &mut w,
+            &format!("{prefix}.gated_positional_embedding.tile_embedding.weight"),
+            &[ar_rows, tiles * np * h],
+            9,
+        );
+        for (name, seed) in [
+            ("pre_tile_positional_embedding", 10),
+            ("post_tile_positional_embedding", 12),
+        ] {
+            put(
+                &mut w,
+                &format!("{prefix}.{name}.embedding.weight"),
+                &[ar_rows, tiles * h],
+                seed,
+            );
+            put(&mut w, &format!("{prefix}.{name}.gate"), &[1], seed + 1);
+        }
+
+        let mut add_layers = |stack: &str, count: usize, gated: bool, base: usize| {
+            for i in 0..count {
+                let p = format!("{prefix}.{stack}.layers.{i}");
+                let s = base + i * 20;
+                put(&mut w, &format!("{p}.input_layernorm.weight"), &[h], s);
+                put(&mut w, &format!("{p}.input_layernorm.bias"), &[h], s + 1);
+                put(
+                    &mut w,
+                    &format!("{p}.post_attention_layernorm.weight"),
+                    &[h],
+                    s + 2,
+                );
+                put(
+                    &mut w,
+                    &format!("{p}.post_attention_layernorm.bias"),
+                    &[h],
+                    s + 3,
+                );
+                for (j, proj) in ["q_proj", "k_proj", "v_proj", "o_proj"].iter().enumerate() {
+                    put(
+                        &mut w,
+                        &format!("{p}.self_attn.{proj}.weight"),
+                        &[h, h],
+                        s + 4 + j,
+                    );
+                }
+                put(&mut w, &format!("{p}.mlp.fc1.weight"), &[inter, h], s + 8);
+                put(&mut w, &format!("{p}.mlp.fc1.bias"), &[inter], s + 9);
+                put(&mut w, &format!("{p}.mlp.fc2.weight"), &[h, inter], s + 10);
+                put(&mut w, &format!("{p}.mlp.fc2.bias"), &[h], s + 11);
+                if gated {
+                    put(&mut w, &format!("{p}.gate_attn"), &[1], s + 12);
+                    put(&mut w, &format!("{p}.gate_ffn"), &[1], s + 13);
+                }
+            }
+        };
+        add_layers("transformer", 2, false, 100);
+        add_layers("global_transformer", 1, true, 200);
+        w
+    }
+
+    fn forward_tower() -> MllamaVisionModel {
+        let config = forward_vision_config();
+        let weights = build_forward_weights("vision_tower");
+        MllamaVisionModel::from_weights(&weights, &config, "vision_tower")
+            .expect("forward-capable tiny tower must load")
+    }
+
+    /// Build `[1, 1, 4, 3, 4, 4]` pixel values: tile 0 is a real tile, tiles
+    /// 1..4 carry the given per-tile fills (zero = processor padding tile).
+    fn pixel_values_with_tiles(tile_fills: [Option<usize>; 4]) -> UniquePtr<MlxArray> {
+        let per_tile = 3 * 4 * 4;
+        let mut pixels = Vec::with_capacity(4 * per_tile);
+        for fill_seed in tile_fills {
+            match fill_seed {
+                Some(seed) => pixels.extend(fill(per_tile, seed)),
+                None => pixels.extend(std::iter::repeat_n(0.0f32, per_tile)),
+            }
+        }
+        mlxcel_core::from_slice_f32(&pixels, &[1, 1, 4, 3, 4, 4])
+    }
+
+    fn max_abs_diff(a: &mlxcel_core::MlxArray, b: &mlxcel_core::MlxArray) -> f32 {
+        let diff = mlxcel_core::subtract(a, b);
+        let m = mlxcel_core::max_all(&mlxcel_core::abs(&diff));
+        mlxcel_core::eval(&m);
+        mlxcel_core::item_f32(&m)
+    }
+
+    /// The content of a padding-tile lane reaches the REAL tile's tower
+    /// output. Two runs share an identical real tile 0 and differ only in
+    /// tile 1 (zero vs non-zero content); the real tile's output changes, so
+    /// padding lanes are live inputs to real-tile results, and slicing them
+    /// away before the tower is NOT an output-preserving optimization.
+    ///
+    /// If this test ever starts failing (diff == 0), the mask semantics were
+    /// changed to truly exclude padding tiles, and a real-tiles-only tower
+    /// fast path would become legal. Revisit issue #527's perf follow-up then.
+    #[test]
+    fn padding_tile_content_reaches_real_tile_output() {
+        let tower = forward_tower();
+        let ar_ids = mlxcel_core::from_slice_i32(&[1], &[1, 1]);
+        let ar_mask = mlxcel_core::from_slice_i32(&[1, 0, 0, 0], &[1, 1, 4]);
+
+        let zero_padding = pixel_values_with_tiles([Some(50), None, None, None]);
+        let perturbed_padding = pixel_values_with_tiles([Some(50), Some(60), None, None]);
+
+        let out_zero = tower.forward(&zero_padding, &ar_ids, &ar_mask);
+        let out_perturbed = tower.forward(&perturbed_padding, &ar_ids, &ar_mask);
+
+        // [1, 1, 4, 5, 16]: hidden (8) + one intermediate layer (8).
+        assert_eq!(
+            mlxcel_core::array_shape(&out_zero),
+            vec![1, 1, 4, 5, 16],
+            "tiny tower output contract"
+        );
+
+        // Compare ONLY the real tile (tile 0).
+        let real_zero = mlxcel_core::slice(&out_zero, &[0, 0, 0, 0, 0], &[1, 1, 1, 5, 16]);
+        let real_perturbed =
+            mlxcel_core::slice(&out_perturbed, &[0, 0, 0, 0, 0], &[1, 1, 1, 5, 16]);
+        let diff = max_abs_diff(&real_zero, &real_perturbed);
+        assert!(
+            diff > 1e-6,
+            "padding-tile content must leak into the real tile's output through \
+             the unmasked real->padding attention (measured diff {diff}); if this \
+             is now 0, the mask semantics changed and a real-tiles-only tower \
+             fast path may have become legal"
+        );
+    }
 }