Fix batched multimodal prefill rank mismatch in Gemma 4 bidirectional mask

[gmhelmold]

Fixes #2323 (Part A)

Root cause

Gemma4TextModel::apply_image_bidirectional_mask assumed batch size 1: it did input_ids.squeeze(0)? followed by to_vec1(), which errors with unexpected rank, expected: 1, got: 2 ([2, 1080]) the moment two same-length requests are batched into a single prefill (input_ids is [bsz, seq]). This path is always taken on cuda-only / no-flash-attn builds. The bidirectional vision-group override was also built as one shared 2-D mask, so even past the crash it couldn't represent per-row image positions.

What changed

Compute the vision-group override per row against the shared 2-D causal mask, then stack the rows into [bsz, 1, seq, total]. bsz == 1 returns the original [seq, total] tensor unchanged, so the single-request path is byte-identical (same dims, dtype, values). [bsz, 1, seq, total] broadcasts to [bsz, n_heads, seq, total] at every consumer of this mask — naive SDPA (broadcast_add), the cuBLASLt rank-4 branch, Metal SDPA (broadcast_as), and the CPU flash-attn mask indexer — for both the full and sliding-window masks.

Tests

Three CPU tests added in vision_models/gemma4/text.rs:

• batched rows equal the independently-computed per-row masks,
• bsz == 1 keeps the old [seq, total] shape and exact values,
• mixed image+video, image-only and text-only rows group correctly (text-only keeps the causal mask untouched).

cargo test -p mistralrs-core --lib gemma4 → 47 passed; 0 failed.
cargo fmt -- --check clean.

Scope

This is Part A (the rank-mismatch crash) only. Part B — mixed-length image batches silently corrupting output — is a separate chunk-coordinate issue and is intentionally left out of this PR.

Happy to have @achandi validate on the deterministic repro mentioned in the issue.

[github-actions]

Code Metrics Report

━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Language              Files        Lines         Code     Comments       Blanks
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 C Header                 23         4454         3116          790          548
 CSS                       3          282          252            6           24
 CUDA                    119        23681        19217         1704         2760
 Dockerfile                1           35           19            9            7
 HTML                      2           27           27            0            0
 JavaScript                3          577          562           12            3
 Jinja2                    7          694          656            5           33
 JSON                     27        15864        15861            0            3
 Makefile                  1           18           16            0            2
 MDX                      36         5901            0         4327         1574
 Metal Shading Lan|       37        14287        11284         1136         1867
 PowerShell                1          656          570           31           55
 Python                  151        12284        10191          484         1609
 Shell                     3         1062          843          117          102
 Plain Text               53        10687            0         9209         1478
 TOML                     28         1368         1189           39          140
 TypeScript               11         1641         1404           66          171
 YAML                      3           25           23            2            0
─────────────────────────────────────────────────────────────────────────────────
 Jupyter Notebooks         3          122           83           23           16
 |- Markdown               1           60           30           22            8
 |- Python                 1          122          113            1            8
 (Total)                              304          226           46           32
─────────────────────────────────────────────────────────────────────────────────
 Markdown                273        11726            0         8736         2990
 |- BASH                  23          295          217           46           32
 |- Dockerfile             2            5            5            0            0
 |- JSON                   6          289          289            0            0
 |- PowerShell             1            1            1            0            0
 |- Python               135         7239         6021          310          908
 |- Rust                  62         3820         2838          388          594
 |- TOML                   7           77           65            0           12
 (Total)                            23452         9436         9480         4536
─────────────────────────────────────────────────────────────────────────────────
 Rust                    674       299702       267081         5872        26749
 |- Markdown             413         9761          452         8126         1183
 (Total)                           309463       267533        13998        27932
─────────────────────────────────────────────────────────────────────────────────
 Svelte                   19         1969         1826           51           92
 |- CSS                    1            4            4            0            0
 |- JavaScript            19          921          767           25          129
 (Total)                             2894         2597           76          221
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
 Total                  1478       429656       345022        41537        43097
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━

[achandi]

We hit this same crash while investigating #2323 and can offer on-hardware validation: we'd been running an equivalent local patch (same shape — per-row override, [bs, 1, seq, total] mask) on real workloads for the past two days.

Setup: DGX Spark (GB10, sm_121, CUDA 13.0), Gemma-4-26B-A4B, Q6K UQFF, 0.8.23 base.

Crash (Part A): gone.

• The deterministic repro from Batched multimodal prefill crashes with rank mismatch; mixed-length image batches silently corrupt output (Gemma 4) #2323 (N identical concurrent image requests, greedy): unexpected rank, expected: 1, got: 2 ([2, 1080]) no longer reproduces at N=2 or N=5.
• Verified on two builds: cuda-only (which always takes this masked path) and cuda + flash-attn.

End-to-end effect: with the equivalent patch, a width-2 batched vision workload (5-page document review) went from near-total output degeneration (every page hitting the token cap with degenerate text) to substantively correct output on our eval — so the fix restores real quality at batch width 2, not just crash-avoidance.

Still broken after the fix (consistent with this PR's "Part A only" scope):

• Mixed-length image batches still corrupt stochastically: individual requests collapse to a handful of output tokens (observed out=8-67 where several hundred are expected), with probability rising with batch width; at width 5 the output is unusable. That's Part B of Batched multimodal prefill crashes with rank mismatch; mixed-length image batches silently corrupt output (Gemma 4) #2323 (chunked-prefill / mixed-length, similar shape to what fix(vlm): handle Qwen VL chunked paged prefill #2301 fixed for Qwen).
• Same-length batched requests complete post-fix, but identical greedy inputs still diverge token-wise (outputs stay coherent) — looks like cross-sequence reduction-order nondeterminism rather than anything in this mask.

One note for review: downstream mask chunking dispatches on rank, so it's worth confirming the always-4D mask preserves single-request behavior — our local variant conservatively kept the 2-D [seq, total] shape for bs == 1; the unit tests added here may already cover this.

Happy to re-run the repro or a width ladder on this exact branch if useful.

[gmhelmold]

Thank you for the on-hardware validation — that's exactly the coverage this needed! Great to hear the deterministic repro is gone at both N=2 and N=5 on GB10, and especially that you verified it on the cuda-only build (which always takes the masked path). With your two-build validation on real Gemma-4-26B workloads plus the byte-parity proof for bs=1 in the PR body, I believe Part A is fully covered here; Part B (batched-prefill rank mismatch) remains scoped out as discussed in #2323.