Onboarding: offline Dia tests + guide updates for quant predicates, KV cache sizing, and batch contracts

docs/contributing/adding-a-model.md

@@ -66,6 +66,15 @@ If your model needs a backbone that is not vendored yet, copy it from mlx-lm
 into `mlx_audio/lm/models/` verbatim and add the provenance header used by the
 other files there (upstream path, version, commit).
 
+!!! warning "Size KV caches for prefill *plus* generation"
+    When your generation loop writes into a fixed-size KV cache, allocate it
+    for the prompt/prefill tokens **and** the generated tokens combined --
+    not just `max_tokens`. If an audio (or text) prompt is prefilled first,
+    those slots are consumed before decoding starts, and undersized caches
+    fail with an index/assertion error mid-generation. A regression test that
+    runs one generation step past the prefill length catches this class of
+    bug; see `mlx_audio/tts/tests/test_dia.py`.
+
 ### Model Configuration
 
 Create a dataclass for your model's config that extends `BaseModelArgs`:
@@ -195,6 +204,37 @@ python -m mlx_audio.convert \
     --dtype bfloat16
 ```
 
+### Prevent Broken Conversions: `model_quant_predicate`
+
+The converter quantizes every `nn.Linear` (and similar) layer whose last
+dimension is a multiple of 64. That blanket policy is wrong for some
+architectures and produces checkpoints that load fine but generate garbage.
+Define a `model_quant_predicate` classmethod on your model when any of the
+following apply:
+
+- An `nn.Embedding` doubles as the tied output/logit head (quantizing it
+  degrades generation logits).
+- A sensitive head or projection drives discrete codebook sampling
+  (e.g. residual/RVQ heads).
+- Any other path where 4-bit weights are known to be lossy for this
+  architecture.
+
+```python
+class Model(nn.Module):
+    @classmethod
+    def model_quant_predicate(cls, path: str, module) -> bool:
+        # Keep embeddings and the fast residual path in full precision.
+        return (
+            not isinstance(module, nn.Embedding)
+            and "fast_" not in path
+        )
+```
+
+The predicate receives `(path, module)` and returns whether quantization is
+allowed; it is combined with the converter's base requirements, so returning
+`True` never forces quantization of an unsuitable tensor. See
+`qwen3_tts`, `fish_qwen3_omni`, or `spark` for examples.
+
 ### Publish to Hugging Face
 
 If you plan to share the converted model, prefer publishing it on the
@@ -209,18 +249,28 @@ the `mlx-community` org when possible.
 
 ### Test
 
-Write a basic test:
+Prefer offline unit tests that build a tiny randomly-initialized model over
+tests that download weights -- they run in seconds, work in CI without
+network, and still exercise real generation code paths:
 
 ```python
-from mlx_audio.tts.utils import load_model
-
 def test_my_model():
-    model = load_model("path/to/my-model-bf16")
+    model = MyModel(tiny_config())  # small dims, random init, no downloads
     results = list(model.generate("Hello, world!"))
     assert len(results) > 0
     assert results[0].audio.shape[0] > 0
 ```
 
+Stub out any component your model downloads at construction time (codec,
+tokenizer) with a fake that returns correctly-shaped arrays; see
+`mlx_audio/tts/tests/test_dia.py` or `test_qwen3_tts.py` for patterns, and
+`mlx_audio/music/tests/test_minimax_music3.py` for a tiny end-to-end config.
+
+If the model supports batch generation, also verify the result contract:
+**exactly one result per input item**, even when an item produces no audio
+(yield silence rather than skipping it) so streaming consumers never see a
+sequence terminate without a final chunk.
+
 Run it:
 
 ```bash
@@ -241,8 +291,10 @@ pytest mlx_audio/tts/tests/test_my_model.py
 - [ ] `__init__.py` exports `Model` and `ModelConfig`
 - [ ] `generate()` method yields `GenerationResult` objects
 - [ ] Model type registered in `MODEL_REMAPPING` (if needed)
-- [ ] Weights converted to MLX `.safetensors` format
+- [ ] KV caches sized for prefill + generation (if the model manages caches manually)
+- [ ] `model_quant_predicate` defined if blanket quantization would corrupt the model
+- [ ] Weights converted to MLX `.safetensors` format and verified after conversion
 - [ ] Hugging Face repo chosen and linked in docs (`mlx-community/...` preferred)
-- [ ] Basic test written and passing
+- [ ] Offline unit test written and passing
 - [ ] Documentation page added
 - [ ] PR submitted with a clear description

mlx_audio/tts/tests/test_dia.py

@@ -0,0 +1,154 @@
+# Copyright (c) 2025, Prince Canuma and contributors (https://github.com/Blaizzy/mlx-audio)
+
+"""Offline tests for the Dia TTS model.
+
+These tests run on a tiny randomly-initialized model so no network access or
+pretrained weights are required. Generation with an audio prompt (ref_audio)
+has additional coverage for the KV cache sizing contract in the PR that
+introduced it; the ``TestKVCacheCapacity`` test below documents that contract.
+"""
+
+import unittest
+from unittest.mock import patch
+
+import mlx.core as mx
+
+from mlx_audio.tts.models.dia.config import (
+    DataConfig,
+    DecoderConfig,
+    EncoderConfig,
+    ModelConfig,
+    TrainingConfig,
+)
+from mlx_audio.tts.models.dia.dia import Model
+
+
+def _tiny_config_dict() -> dict:
+    return {
+        "version": "1.0",
+        "model": {
+            "encoder": {
+                "n_layer": 1,
+                "n_embd": 16,
+                "n_hidden": 32,
+                "n_head": 2,
+                "head_dim": 8,
+            },
+            "decoder": {
+                "n_layer": 1,
+                "n_embd": 16,
+                "n_hidden": 32,
+                "gqa_query_heads": 2,
+                "kv_heads": 1,
+                "gqa_head_dim": 8,
+                "cross_query_heads": 2,
+                "cross_head_dim": 8,
+            },
+            "src_vocab_size": 128,
+            "tgt_vocab_size": 1028,
+            "sample_rate": 44100,
+        },
+        "training": {"dtype": "float32", "logits_dot_in_fp32": False},
+        "data": {
+            "text_length": 128,
+            "audio_length": 256,
+            "channels": 9,
+            "audio_eos_value": 1024,
+            "audio_pad_value": 1025,
+            "audio_bos_value": 1026,
+            "delay_pattern": [0, 8, 9, 10, 11, 12, 13, 14, 15],
+        },
+    }
+
+
+class FakeQuantizer:
+    def from_codes(self, audio_codes):
+        codes = mx.array(audio_codes, dtype=mx.float32)
+        return (codes,)
+
+
+class FakeDAC:
+    """Stands in for the DAC codec so tests never touch the network."""
+
+    def __init__(self):
+        self.quantizer = FakeQuantizer()
+
+    @classmethod
+    def from_pretrained(cls, repo_id):
+        return cls()
+
+    def preprocess(self, input_values, sample_rate):
+        return input_values
+
+    def encode(self, audio_data, n_quantizers=None):
+        # One frame of valid codebook indices per channel: shape (1, C, T).
+        frame = mx.zeros((1, 9, 8), dtype=mx.int32)
+        return None, frame, None, None, None
+
+    def decode(self, audio_values):
+        return mx.zeros((1, 1, 64), dtype=mx.float32)
+
+
+def _make_model() -> Model:
+    with patch("mlx_audio.tts.models.dia.dia.DAC", FakeDAC):
+        return Model(_tiny_config_dict())
+
+
+class TestDiaConfigParsing(unittest.TestCase):
+    def test_tiny_config_loads(self):
+        config_dict = _tiny_config_dict()
+        model = _make_model()
+        self.assertEqual(model.config.model.encoder.n_layer, 1)
+        self.assertEqual(model.config.data.channels, 9)
+        self.assertIsInstance(config_dict["model"]["decoder"]["kv_heads"], int)
+
+    def test_model_config_dataclasses_roundtrip(self):
+        encoder = EncoderConfig(n_layer=1, n_embd=16, n_hidden=32, n_head=2, head_dim=8)
+        decoder = DecoderConfig(
+            n_layer=1,
+            n_embd=16,
+            n_hidden=32,
+            gqa_query_heads=2,
+            kv_heads=1,
+            gqa_head_dim=8,
+            cross_query_heads=2,
+            cross_head_dim=8,
+        )
+        model_config = ModelConfig(encoder=encoder, decoder=decoder)
+        training = TrainingConfig()
+        data = DataConfig(text_length=128, audio_length=256)
+        # Lengths are rounded up to multiples of 128.
+        self.assertEqual(data.audio_length % 128, 0)
+        self.assertEqual(data.text_length % 128, 0)
+        self.assertEqual(training.dtype, "bfloat16")
+        self.assertIsNotNone(model_config)
+
+
+class TestDiaGeneration(unittest.TestCase):
+    def test_generation_without_ref_audio_completes(self):
+        model = _make_model()
+        result = next(
+            model.generate("Hello world", max_tokens=32, verbose=False),
+        )
+        self.assertIsInstance(result.audio, mx.array)
+        self.assertGreater(result.audio.shape[0], 0)
+        self.assertEqual(result.sample_rate, 44100)
+
+    def test_kv_cache_rejects_writes_beyond_capacity(self):
+        # Documents the failure mode the ref_audio path used to hit when the
+        # cache was sized only for the generated tokens.
+        from mlx_audio.tts.models.dia.layers import KVCache
+
+        cache = KVCache(num_heads=2, max_len=3, head_dim=4)
+        cache.prefill_kv(
+            mx.zeros((2, 2, 2, 4)), mx.zeros((2, 2, 2, 4))
+        )  # prefill fills slots 0..1
+        cache.update_and_fetch(
+            mx.zeros((2, 2, 1, 4)), mx.zeros((2, 2, 1, 4))
+        )  # writes slot 2
+        with self.assertRaises(AssertionError):
+            cache.update_and_fetch(mx.zeros((2, 2, 1, 4)), mx.zeros((2, 2, 1, 4)))
+
+
+if __name__ == "__main__":
+    unittest.main()