SIGN IN SIGN UP

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