fix: replace NaN from all-masked SDPA padding rows in Gemma 4 vision

[fabiopili]

I've been having problems with Gemma 4 models that I quantized locally and enabled the full 1120 vision token budget when serving from oMLX and accessing from Open WebUI. The model would respond an endless loop of the <pad> token.

A code review is recommended before merging, even though the impact seems minimal. I'm not a proficient Python developer and this change was created with Claude Code assistance for the bug identification. That said, I've been using it locally without problems and tests pass.

---

Summary

• Replace NaN values produced by softmax over all-masked padding rows in Gemma 4 vision attention with 0, preventing NaN propagation through residual connections and layer norms
• The existing ensure_fused_sdpa workaround pads head_dim to force the fused SDPA kernel, but does not help at sequence lengths (e.g. L=10080 for 1120-token budget) where the non-fused fallback is used
• Padding rows are zeroed out by VisionPooler before pooling anyway, so 0 is the correct replacement value

Problem

MLX's scaled_dot_product_attention computes softmax(QK^T) V. When an entire row of the attention mask is -inf (padding tokens that should attend to nothing), softmax produces 0/0 = NaN. This NaN then propagates through:

1. The output projection (o_proj)
2. Residual additions in each of the 24 transformer layers
3. RMSNorm (any NaN input produces NaN output)
4. The MLP branch
5. All subsequent layers

By the time the VisionPooler zeros out padding positions, the NaN has already corrupted non-padding rows via residual connections.

Fix

A single line after the SDPA call in VisionAttention.__call__:

attn_output = mx.where(mx.isnan(attn_output), 0.0, attn_output)

This seems safe because:

• Only padding rows produce NaN (real content rows have valid attention targets)
• VisionPooler explicitly zeros padding positions before pooling (vision.py:364-367)
• The replacement runs inside @mx.compile scope, so the isnan check fuses into the existing compute graph with minimal overhead

Changed files

• mlx_vlm/models/gemma4/vision.py (+5 lines)

Test plan

• All 395 unit tests pass (0 failures, 1 skipped, 2 unrelated errors)
• End-to-end inference with Gemma 4 on single image (fused SDPA path)
• End-to-end inference with Gemma 4 on multiple different-sized images (non-fused path, L=10080)
• Verify no NaN in output embeddings at both sequence lengths