Repository: ml-explore/mlx-lm

Issue:
gemma-3-27b-it-qat is not working with KV cache quantization
mlx-lm version: `dd24e000b6c002e7bf0ffd88295db0a0e1008e88`
mlx version: `0.25.0`

Commands to reproduce:
```bash
python -m mlx_lm convert --hf-path /Volumes/T7/Models/hf-models/gemma-3-27b-it-qat-q4_0-unquantized --mlx-path gemma-3-27b-it-qat-4bit-text --dtype bfloat16 -q --q-bits 4
cat very_long_prompt.txt | python -m mlx_lm generate --model gemma-3-27b-it-qat-4bit-text --prompt - --kv-bits 8 --kv-group-size 64 --quantized-kv-start 5000
```

Error message:
```bash
Traceback (most recent call last):
  File "<frozen runpy>", line 198, in _run_module_as_main
  File "<frozen runpy>", line 88, in _run_code
  File "/Users/neil/workspace/mlx-lm/mlx_lm/__main__.py", line 27, in <module>
    submodule.main()
  File "/Users/neil/workspace/mlx-lm/mlx_lm/generate.py", line 814, in main
    response = generate(
               ^^^^^^^^^
  File "/Users/neil/workspace/mlx-lm/mlx_lm/generate.py", line 689, in generate
    for response in stream_generate(model, tokenizer, prompt, **kwargs):
  File "/Users/neil/workspace/mlx-lm/mlx_lm/generate.py", line 622, in stream_generate
    for n, (token, logprobs, from_draft) in enumerate(token_generator):
  File "/Users/neil/workspace/mlx-lm/mlx_lm/generate.py", line 611, in <genexpr>
    token_generator = (
                      ^
  File "/Users/neil/workspace/mlx-lm/mlx_lm/generate.py", line 363, in generate_step
    model(y[:prefill_step_size][None], cache=prompt_cache)
  File "/Users/neil/workspace/mlx-lm/mlx_lm/models/gemma3.py", line 45, in __call__
    return self.language_model(inputs, cache=cache, mask=mask)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/neil/workspace/mlx-lm/mlx_lm/models/gemma3_text.py", line 223, in __call__
    out = self.model(inputs, mask, cache)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/neil/workspace/mlx-lm/mlx_lm/models/gemma3_text.py", line 204, in __call__
    h = layer(h, local_mask, c)
        ^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/neil/workspace/mlx-lm/mlx_lm/models/gemma3_text.py", line 153, in __call__
    r = self.self_attn(self.input_layernorm(x), mask, cache)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/neil/workspace/mlx-lm/mlx_lm/models/gemma3_text.py", line 93, in __call__
    output = mx.fast.scaled_dot_product_attention(
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: scaled_dot_product_attention(): incompatible function arguments. The following argument types are supported:
    1. scaled_dot_product_attention(q: array, k: array, v: array, *, scale: float,  mask: Union[None, str, array] = None, stream: Union[None, Stream, Device] = None) -> array

Invoked with types: mlx.core.array, tuple, tuple, kwargs = { scale: float, mask: str }
```

Discussion:

Solution patch:
diff --git a/mlx_lm/models/base.py b/mlx_lm/models/base.py
index fe9d57a..19aa8c6 100644
--- a/mlx_lm/models/base.py
+++ b/mlx_lm/models/base.py
@@ -89,7 +89,15 @@ def quantized_scaled_dot_product_attention(
         queries, *q_keys, transpose=True, group_size=group_size, bits=bits
     )
     if mask is not None:
-        scores += mask
+        if isinstance(mask, str):
+            qL, kL = scores.shape[-2:]
+            q_indices = mx.arange(kL - qL, kL)
+            k_indices = mx.arange(kL)
+            mask = q_indices[:, None] >= k_indices[None]
+        if mask.dtype == mx.bool_:
+            scores = mx.where(mask, scores, mx.finfo(scores.dtype).min)
+        else:
+            scores += mask
     scores = mx.softmax(scores, axis=-1, precise=True)
     out = mx.quantized_matmul(
         scores, *q_values, transpose=False, group_size=group_size, bits=bits
diff --git a/mlx_lm/models/gemma3_text.py b/mlx_lm/models/gemma3_text.py
index 9d00263..fdcc278 100644
--- a/mlx_lm/models/gemma3_text.py
+++ b/mlx_lm/models/gemma3_text.py
@@ -7,7 +7,7 @@ from typing import Any, Optional
 import mlx.core as mx
 import mlx.nn as nn
 
-from .base import BaseModelArgs, create_attention_mask
+from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .cache import KVCache, RotatingKVCache
 
 
@@ -89,9 +89,8 @@ class Attention(nn.Module):
         # Sliding window
         if isinstance(mask, mx.array) and mask.shape[-1] != keys.shape[-2]:
             mask = mask[..., -keys.shape[-2] :]
-
-        output = mx.fast.scaled_dot_product_attention(
-            queries, keys, values, scale=self.scale, mask=mask
+        output = scaled_dot_product_attention(
+            queries, keys, values, cache=cache, scale=self.scale, mask=mask
         )
         output = output.transpose(0, 2, 1, 3).reshape(B, L, -1)
         return self.o_proj(output)


Test patch:
diff --git a/tests/test_models.py b/tests/test_models.py
index 163263b..ae989ed 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -6,7 +6,7 @@ import mlx.nn as nn
 from mlx.utils import tree_map
 
 from mlx_lm.models import rope_utils
-from mlx_lm.models.base import create_causal_mask
+from mlx_lm.models.base import create_causal_mask, scaled_dot_product_attention
 from mlx_lm.models.cache import KVCache, RotatingKVCache, make_prompt_cache
 
 
@@ -166,6 +166,42 @@ class TestModels(unittest.TestCase):
         )
         self.assertTrue(isinstance(rope, rope_utils.Llama3RoPE))
 
+    def test_quantized_sdpa(self):
+        cache = KVCache()
+
+        k = 1e-1 * mx.random.normal(shape=(1, 1, 256, 32))
+        v = 1e-1 * mx.random.normal(shape=(1, 1, 256, 32))
+
+        cache.update_and_fetch(k, v)
+        quant_cache = cache.to_quantized(group_size=32, bits=8)
+
+        k = 1e-1 * mx.random.normal(shape=(1, 1, 1, 32))
+        v = 1e-1 * mx.random.normal(shape=(1, 1, 1, 32))
+
+        k_up, v_up = cache.update_and_fetch(k, v)
+        qk_up, qv_up = quant_cache.update_and_fetch(k, v)
+
+        q = 1e-1 * mx.random.normal(shape=(1, 4, 257, 32))
+
+        mask = "causal"
+        out = scaled_dot_product_attention(
+            q,
+            k_up,
+            v_up,
+            cache=cache,
+            mask=mask,
+            scale=1.0,
+        )
+        qout = scaled_dot_product_attention(
+            q,
+            qk_up,
+            qv_up,
+            cache=quant_cache,
+            mask=mask,
+            scale=1.0,
+        )
+        self.assertTrue(mx.allclose(out, qout, rtol=1e-2, atol=1e-2))
+
     def model_test_runner(self, model, model_type, vocab_size, num_layers):
 
         self.assertEqual(len(model.layers), num_layers)