fix(attention): materialize a quantized KV state before hand-rolled attention (#2065)
gemma2 (logit softcapping) and phi3small (block-sparse mask) build their attention scores by hand instead of going through scaled_dot_product_attention. QuantizedKVCache.update_and_fetch returns (packed, scales, biases) tuples, so mx.expand_dims raises TypeError: expand_dims(): incompatible function arguments as soon as a quantized KV cache reaches either path. Add dequantize_kv_state() next to the other KV-state helpers in base.py and call it on the two hand-rolled paths. Fixes #2060 Co-authored-by: Alazer Manakelew <alazermanakelewb@gmail.com>
T
Tai An committed
9606b86d9746962eaa505064901b66aef1a6df13
Parent: 8adff01
Committed by GitHub <noreply@github.com>
on 8/28/2026, 1:59:04 PM