Fix broadcast crash in quantized SDPA with GQA + batched padding mask (batch >= 2)

[pinglin]

Summary

quantized_scaled_dot_product_attention crashes with [broadcast_shapes] … cannot be broadcast for any grouped-query-attention model using a quantized KV cache once the batch size is ≥ 2. This makes kv-bits unusable with batched decoding on GQA models (e.g. Qwen3, Qwen3.5/3.6). Tested on main @ 2ed22318.

Root cause

For GQA (n_repeats = n_q_heads // n_kv_heads > 1) the function reshapes queries to 5D (B, n_kv_heads, n_repeats, L, D), so scores become 5D (B, n_kv_heads, n_repeats, L, L) — but the mask is passed through unchanged. A batched padding mask is 4D (B, 1, L, L). Right-aligned broadcasting lines the mask's B axis up against n_kv_heads: harmless at B == 1 (1-vs-n_kv_heads), but at B ≥ 2 it becomes B-vs-n_kv_heads (e.g. 2 vs 8) with neither equal to 1 → broadcast failure.

Minimal repro

Direct call to quantized_scaled_dot_product_attention with GQA dims (n_q_heads=40, n_kv_heads=8 → n_repeats=5), 8-bit quantized KV, batch=2, a (2,1,L,L) bool mask:

queries (2, 40, 16, 128)  q_keys[0] (2, 8, 16, 32)  mask (2, 1, 16, 16)  n_repeats=5
ValueError: [broadcast_shapes] Shapes (2,1,16,16) and (2,8,5,16,16) cannot be broadcast.
  at mlx_lm/models/base.py:94  scores = mx.where(mask, scores, mx.finfo(scores.dtype).min)

Fix

             mask = q_indices[:, None] >= k_indices[None]
+        if n_repeats > 1 and mask.ndim >= 3:
+            mask = mx.expand_dims(mask, -3)
         if mask.dtype == mx.bool_:
             scores = mx.where(mask, scores, mx.finfo(scores.dtype).min)
         else:
             scores += mask

Inserts a singleton n_repeats axis so (B,1,L,L) → (B,1,1,L,L) and broadcasts across the repeat groups. Guarded on mask.ndim >= 3 so the 2D causal-string path and the MHA (n_repeats == 1) path are untouched.

Verification

Patched, same B=2 GQA quantized-KV inputs:

OK: out.shape=(2, 40, 16, 128) finite=True
row1 padded-column attention-weight sum=0.000e+00   # mask applied correctly, not just made to broadcast

Also verified end-to-end: Qwen3.6-35B-A3B with --kv-bits 8 and 8-way in-flight — which crashed on every batch ≥ 2 before — now serves concurrent generations cleanly (0 broadcast errors, coherent output).