mlx-lm/tests/test_models.py

# Copyright © 2024 Apple Inc.
import copy
import importlib
import unittest

import mlx.core as mx
import mlx.nn as nn
from mlx.utils import tree_flatten, tree_map

from mlx_lm.models import rope_utils
from mlx_lm.models.base import create_causal_mask, scaled_dot_product_attention
from mlx_lm.models.cache import KVCache, RotatingKVCache, make_prompt_cache
from mlx_lm.models.gated_delta import (
    gated_delta_kernel,
    gated_delta_ops,
    gated_delta_update,
)
from mlx_lm.models.ssm import ssm_attn, ssm_update


class TestModels(unittest.TestCase):
    def test_kv_cache(self):
        cache = KVCache()

        k = mx.ones((1, 4, 1, 32), mx.float16)
        v = mx.ones((1, 4, 1, 32), mx.float16)

        k_up, v_up = cache.update_and_fetch(k, v)
        self.assertTrue(mx.array_equal(k_up, k))
        self.assertTrue(mx.array_equal(v_up, v))
        self.assertEqual(cache.offset, 1)

        k = mx.ones((1, 4, cache.step, 32), mx.float16)
        v = mx.ones((1, 4, cache.step, 32), mx.float16)
        k_up, v_up = cache.update_and_fetch(k, v)

        expected = mx.ones((1, 4, cache.step + 1, 32), mx.float16)
        self.assertTrue(mx.array_equal(k_up, expected))
        self.assertTrue(mx.array_equal(v_up, expected))
        self.assertEqual(cache.offset, cache.step + 1)

    def test_rotating_kv_cache(self):
        b, h, d = 1, 2, 32
        cache = RotatingKVCache(max_size=8)

        k = mx.random.uniform(shape=(b, h, 2, d))
        v = mx.random.uniform(shape=(b, h, 2, d))

        k_up, v_up = cache.update_and_fetch(k, v)
        self.assertTrue(mx.array_equal(k_up, k))
        self.assertTrue(mx.array_equal(v_up, v))
        self.assertEqual(cache.offset, 2)

        k = mx.random.uniform(shape=(b, h, 5, d))
        v = mx.random.uniform(shape=(b, h, 5, d))
        k_up, v_up = cache.update_and_fetch(k, v)
        self.assertTrue(mx.array_equal(k_up[..., 2:, :], k))
        self.assertTrue(mx.array_equal(v_up[..., 2:, :], v))

        k = mx.random.uniform(shape=(b, h, 4, d))
        v = mx.random.uniform(shape=(b, h, 4, d))
        k_up, v_up = cache.update_and_fetch(k, v)
        self.assertTrue(mx.array_equal(k_up[..., -4:, :], k))
        self.assertTrue(mx.array_equal(v_up[..., -4:, :], v))

        idx = 0
        for _ in range(10):
            k = mx.random.uniform(shape=(b, h, 1, d))
            v = mx.random.uniform(shape=(b, h, 1, d))
            k_up, v_up = cache.update_and_fetch(k, v)
            self.assertTrue(mx.array_equal(k_up[..., idx : idx + 1, :], k))
            self.assertTrue(mx.array_equal(v_up[..., idx : idx + 1, :], v))
            idx += 1
            idx %= 8

        # Try with nonzero keep
        cache = RotatingKVCache(max_size=8, keep=2)

        # Check a large update
        k = mx.random.uniform(shape=(b, h, 20, d))
        v = mx.random.uniform(shape=(b, h, 20, d))
        k_up, v_up = cache.update_and_fetch(k, v)
        self.assertTrue(mx.array_equal(k_up, k))
        self.assertTrue(mx.array_equal(v_up, v))

        # A bunch of small updates
        self.assertEqual(cache.offset, 20)
        idx = 2
        for i in range(10):
            k = mx.random.uniform(shape=(b, h, 1, d))
            v = mx.random.uniform(shape=(b, h, 1, d))
            k_up, v_up = cache.update_and_fetch(k, v)
            self.assertTrue(mx.array_equal(k_up[..., idx : idx + 1, :], k))
            self.assertTrue(mx.array_equal(v_up[..., idx : idx + 1, :], v))
            self.assertEqual(cache.offset, 21 + i)
            idx += 1
            if idx >= 8:
                idx = 2

    def test_rotating_kv_cache_chat_mode(self):
        # Test that the rotating kv cache can handle
        # alternating prompt/prefill with generation
        d = 4
        h = 2
        cache = RotatingKVCache(max_size=18)

        x = mx.random.uniform(shape=(1, h, 8, d))
        k, v = cache.update_and_fetch(x, x)
        self.assertEqual(k.shape[2], 8)
        self.assertEqual(cache.offset, 8)

        x = mx.random.uniform(shape=(1, h, 1, d))
        k, v = cache.update_and_fetch(x, x)
        self.assertEqual(k.shape[2], 9)
        self.assertEqual(cache.offset, 9)
        self.assertTrue(mx.allclose(x, k[..., 8:9, :]))

        x = mx.random.uniform(shape=(1, h, 2, d))
        k, v = cache.update_and_fetch(x, x)
        self.assertEqual(k.shape[2], 11)
        self.assertEqual(cache.offset, 11)
        self.assertTrue(mx.allclose(x, k[..., 9:11, :]))

        x = mx.random.uniform(shape=(1, h, 3, d))
        k, v = cache.update_and_fetch(x, x)
        self.assertEqual(k.shape[2], 14)
        self.assertEqual(cache.offset, 14)
        self.assertTrue(mx.allclose(x, k[..., 11:14, :]))

        x = mx.random.uniform(shape=(1, h, 6, d))
        k, v = cache.update_and_fetch(x, x)
        self.assertEqual(cache.offset, 20)
        self.assertTrue(mx.allclose(x, k[..., -6:, :]))

        x = mx.random.uniform(shape=(1, h, 2, d))
        k, v = cache.update_and_fetch(x, x)
        self.assertEqual(cache.offset, 22)
        self.assertTrue(mx.allclose(x, k[..., -2:, :]))

    def test_causal_mask_padding(self):
        right_padding = mx.array([2, 1, 0])
        mask = create_causal_mask(3, right_padding=right_padding)

        causal_mask = create_causal_mask(3)
        self.assertTrue(
            mx.array_equal(mask[0, 0], causal_mask & mx.array([True, False, False]))
        )
        self.assertTrue(
            mx.array_equal(mask[1, 0], causal_mask & mx.array([True, True, False]))
        )
        self.assertTrue(mx.array_equal(mask[2, 0], causal_mask))

        left_padding = mx.array([2, 1, 0])
        mask = create_causal_mask(3, left_padding=left_padding)

        self.assertTrue(
            mx.array_equal(mask[0, 0], causal_mask & mx.array([False, False, True]))
        )
        self.assertTrue(
            mx.array_equal(mask[1, 0], causal_mask & mx.array([False, True, True]))
        )
        self.assertTrue(mx.array_equal(mask[2, 0], causal_mask))

    def test_mask_with_window(self):
        mask = create_causal_mask(5, 0, window_size=3)
        expected_sums = mx.array([1, 2, 3, 3, 3])
        sums = mask.sum(axis=1)
        self.assertTrue(mx.array_equal(sums, expected_sums))

        mask = create_causal_mask(5, 1, window_size=3)
        self.assertEqual(mask.shape, (5, 6))
        expected_sums = mx.array([2, 3, 3, 3, 3])
        sums = mask.sum(axis=1)
        self.assertTrue(mx.array_equal(sums, expected_sums))

        mask = create_causal_mask(5, 2, window_size=3)
        self.assertEqual(mask.shape, (5, 7))
        expected_sums = mx.array([3, 3, 3, 3, 3])
        sums = mask.sum(axis=1)
        self.assertTrue(mx.array_equal(sums, expected_sums))

    def test_llama_model_sliding_attention(self):
        from mlx_lm.models import llama

        args = llama.ModelArgs(
            model_type="llama",
            hidden_size=64,
            num_hidden_layers=4,
            intermediate_size=256,
            num_attention_heads=8,
            num_key_value_heads=4,
            rms_norm_eps=1e-5,
            vocab_size=128,
            sliding_window=4,
            layer_types=[
                "full_attention",
                "sliding_attention",
                "sliding_attention",
                "full_attention",
            ],
            tie_word_embeddings=False,
            rope_theta=10000.0,
        )
        model = llama.Model(args)

        tokens = mx.array([[1, 2, 3, 4, 5]], dtype=mx.int32)
        out = model(tokens)
        mx.eval(out)
        self.assertEqual(out.shape, (1, 5, args.vocab_size))

        caches = model.make_cache()
        self.assertIsInstance(caches[0], KVCache)
        self.assertIsInstance(caches[1], RotatingKVCache)
        self.assertIsInstance(caches[2], RotatingKVCache)
        self.assertIsInstance(caches[3], KVCache)

        caches = model.make_cache()
        step = model(tokens[:, :2], cache=caches)
        mx.eval(step)
        step = model(tokens[:, 2:3], cache=caches)
        mx.eval(step)
        self.assertEqual(caches[0].offset, 3)
        self.assertEqual(caches[1].offset, 3)

    def test_rope(self):
        rope = rope_utils.initialize_rope(32, base=100, traditional=False)
        self.assertTrue(isinstance(rope, nn.RoPE))

        rope = rope_utils.initialize_rope(
            32,
            base=100,
            traditional=False,
            scaling_config={"rope_type": "linear", "factor": 10.0},
        )
        self.assertTrue(isinstance(rope, nn.RoPE))

        rope = rope_utils.initialize_rope(
            32,
            base=100,
            traditional=False,
            scaling_config={"rope_type": "llama3", "factor": 2.0},
        )
        self.assertTrue(isinstance(rope, rope_utils.Llama3RoPE))

        rope = rope_utils.initialize_rope(
            16,
            base=100.0,
            traditional=False,
            scaling_config={
                "rope_type": "proportional",
                "partial_rotary_factor": 0.5,
            },
        )
        self.assertTrue(isinstance(rope, rope_utils.ProportionalRoPE))
        expected_freqs = 100.0 ** (mx.arange(0, 8, 2, dtype=mx.float32) / 16)
        self.assertTrue(mx.allclose(rope._freqs[:4], expected_freqs))
        self.assertTrue(mx.all(mx.isinf(rope._freqs[4:])))

        x = mx.arange(16, dtype=mx.float32).reshape(1, 1, 1, 16)
        y = rope(x, offset=1)
        expected_rotated = mx.fast.rope(
            mx.concatenate([x[..., :4], x[..., 8:12]], axis=-1),
            8,
            traditional=False,
            base=None,
            scale=1.0,
            offset=1,
            freqs=expected_freqs,
        )
        expected = mx.concatenate(
            [
                expected_rotated[..., :4],
                x[..., 4:8],
                expected_rotated[..., 4:],
                x[..., 12:],
            ],
            axis=-1,
        )
        mx.eval(y, expected)
        self.assertTrue(mx.allclose(y, expected))

    def test_su_scaled_rope_no_mutation(self):
        rope = rope_utils.SuScaledRoPE(
            dims=8,
            max_position_embeddings=131072,
            original_max_position_embeddings=4096,
            long_factor=[1.0] * 4,
        )
        x = mx.ones((1, 2, 4, 8))
        rope(x)
        mx.eval(x)
        self.assertTrue((x == 1).all())

    def test_yarn_rope_no_mutation(self):
        rope = rope_utils.YarnRoPE(
            dims=8,
            scaling_factor=2.0,
            mscale=1.0,
            mscale_all_dim=0,
        )
        x = mx.ones((1, 2, 4, 8))
        rope(x)
        mx.eval(x)
        self.assertTrue((x == 1).all())

    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)
        self.assertEqual(model.model_type, model_type)

        for t in [mx.float32, mx.float16]:
            model.update(tree_map(lambda p: p.astype(t), model.parameters()))

            inputs = mx.array([[0, 1]])
            outputs = model(inputs)
            self.assertEqual(outputs.shape, (1, 2, vocab_size))
            self.assertEqual(outputs.dtype, t)

            cache = make_prompt_cache(model)
            outputs = model(inputs, cache=cache)
            self.assertEqual(outputs.shape, (1, 2, vocab_size))
            self.assertEqual(outputs.dtype, t)

            outputs = model(mx.argmax(outputs[0, -1:, :], keepdims=True), cache=cache)
            self.assertEqual(outputs.shape, (1, 1, vocab_size))
            self.assertEqual(outputs.dtype, t)

        # Test batch size > 1
        inputs = mx.array([[0, 1], [2, 3]])
        outputs = model(inputs)
        self.assertEqual(outputs.shape, (2, 2, vocab_size))

        # Make sure the model can be copied / pickled
        copy.deepcopy(model)

    def test_llama(self):
        from mlx_lm.models import llama

        args = llama.ModelArgs(
            model_type="llama",
            hidden_size=1024,
            num_hidden_layers=4,
            intermediate_size=2048,
            num_attention_heads=4,
            rms_norm_eps=1e-5,
            vocab_size=10_000,
        )
        model = llama.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_lfm2(self):
        from mlx_lm.models import lfm2

        args = lfm2.ModelArgs(
            model_type="lfm2",
            hidden_size=1024,
            num_hidden_layers=4,
            num_attention_heads=4,
            num_key_value_heads=2,
            norm_eps=1e-5,
            vocab_size=10_000,
            full_attn_idxs=[0, 1, 2],
            rope_theta=10000,
            block_dim=1024,
            block_ffn_dim_multiplier=1.5,
            block_auto_adjust_ff_dim=True,
            block_ff_dim=2048,
            block_multiple_of=256,
            max_position_embeddings=1000,
            conv_bias=True,
            conv_L_cache=3,
        )
        model = lfm2.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_lfm2_moe(self):
        from mlx_lm.models import lfm2_moe

        args = lfm2_moe.ModelArgs(
            model_type="lfm2_moe",
            hidden_size=1024,
            intermediate_size=7168,
            num_hidden_layers=4,
            num_attention_heads=4,
            num_key_value_heads=2,
            norm_eps=1e-5,
            vocab_size=10_000,
            full_attn_idxs=[0, 1, 2],
            rope_theta=10000,
            max_position_embeddings=1000,
            conv_bias=True,
            conv_L_cache=3,
            moe_intermediate_size=1792,
            num_dense_layers=2,
            num_experts=4,
            num_experts_per_tok=2,
            norm_topk_prob=True,
            use_expert_bias=True,
        )
        model = lfm2_moe.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_bitnet(self):
        from mlx_lm.models import bitnet

        args = bitnet.ModelArgs(
            model_type="bitnet",
            hidden_size=1024,
            num_hidden_layers=4,
            intermediate_size=2048,
            num_attention_heads=4,
            num_key_value_heads=4,
            rms_norm_eps=1e-5,
            vocab_size=10_000,
        )
        model = bitnet.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_phi2(self):
        from mlx_lm.models import phi

        args = phi.ModelArgs()
        model = phi.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_phi3(self):
        from mlx_lm.models import phi3

        args = phi3.ModelArgs(
            model_type="phi3",
            hidden_size=3072,
            num_hidden_layers=32,
            intermediate_size=8192,
            num_attention_heads=32,
            rms_norm_eps=1e-5,
            vocab_size=32064,
        )
        model = phi3.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_gemma(self):
        from mlx_lm.models import gemma

        args = gemma.ModelArgs(
            model_type="gemma",
            hidden_size=1024,
            num_hidden_layers=4,
            intermediate_size=2048,
            num_attention_heads=4,
            head_dim=128,
            rms_norm_eps=1e-5,
            vocab_size=10_000,
            num_key_value_heads=4,
        )
        model = gemma.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_mixtral(self):
        from mlx_lm.models import mixtral

        # Make a baby mixtral, because it will actually do the
        # eval
        args = mixtral.ModelArgs(
            model_type="mixtral",
            vocab_size=100,
            hidden_size=32,
            intermediate_size=128,
            num_hidden_layers=2,
            num_attention_heads=4,
            num_experts_per_tok=2,
            num_key_value_heads=2,
            num_local_experts=4,
        )
        model = mixtral.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    @unittest.skip("requires ai2-olmo")
    def test_olmo(self):
        from mlx_lm.models import olmo

        args = olmo.ModelArgs(
            model_type="olmo",
            d_model=1024,
            n_layers=4,
            mlp_hidden_size=2048,
            n_heads=2,
            vocab_size=10_000,
            embedding_size=10_000,
        )
        model = olmo.Model(args)
        self.model_test_runner(
            model,
            args.model_type,
            args.vocab_size,
            args.n_layers,
        )

    def test_qwen3_moe(self):
        from mlx_lm.models import qwen3_moe

        args = qwen3_moe.ModelArgs(
            model_type="qwen3_moe",
            hidden_size=1024,
            num_hidden_layers=4,
            intermediate_size=2048,
            num_attention_heads=4,
            num_key_value_heads=4,
            rms_norm_eps=1e-5,
            head_dim=128,
            vocab_size=10_000,
            decoder_sparse_step=1,
            mlp_only_layers=[],
            num_experts_per_tok=4,
            num_experts=16,
            moe_intermediate_size=1024,
            rope_theta=1000,
            max_position_embeddings=4096,
            tie_word_embeddings=False,
            norm_topk_prob=True,
        )
        model = qwen3_moe.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_qwen3(self):
        from mlx_lm.models import qwen3

        args = qwen3.ModelArgs(
            model_type="qwen3",
            hidden_size=1024,
            num_hidden_layers=4,
            intermediate_size=2048,
            num_attention_heads=4,
            num_key_value_heads=4,
            rms_norm_eps=1e-5,
            vocab_size=10_000,
            head_dim=128,
            max_position_embeddings=4096,
            tie_word_embeddings=False,
            rope_theta=1000,
        )
        model = qwen3.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_qwen3_5_family_convert_then_load_norm_not_shift_twice(self):
        text_config = {
            "hidden_size": 8,
            "intermediate_size": 16,
            "num_hidden_layers": 1,
            "num_attention_heads": 1,
            "num_key_value_heads": 1,
            "rms_norm_eps": 1e-5,
            "vocab_size": 32,
            "linear_num_value_heads": 1,
            "linear_num_key_heads": 1,
            "linear_key_head_dim": 4,
            "linear_value_head_dim": 4,
            "linear_conv_kernel_dim": 1,
            "full_attention_interval": 1,
            "tie_word_embeddings": False,
            "max_position_embeddings": 64,
        }
        hf_norm_key = "model.language_model.layers.0.input_layernorm.weight"
        mlx_norm_key = "language_model.model.layers.0.input_layernorm.weight"

        for model_type, hf_mtp_key in (
            ("qwen3_5", "mtp.fc.weights"),
            ("qwen3_5_moe", "mtp.fc.weight"),
        ):
            module = importlib.import_module(f"mlx_lm.models.{model_type}")
            args = module.ModelArgs.from_dict(
                {
                    "model_type": model_type,
                    "text_config": {"model_type": model_type, **text_config},
                }
            )
            model = module.Model(args)

            base = mx.arange(8, dtype=mx.float32)

            # Simulate convert sanitize on HF-style keys.
            converted = model.sanitize(
                {
                    hf_norm_key: base,
                    hf_mtp_key: mx.zeros((1,), dtype=mx.float32),
                }
            )
            self.assertIn(mlx_norm_key, converted)
            self.assertTrue(mx.array_equal(converted[mlx_norm_key], base + 1.0))
            self.assertFalse(any("mtp." in k for k in converted))

            # Simulate load sanitize on already-converted keys.
            loaded = model.sanitize(converted)
            self.assertTrue(
                mx.array_equal(loaded[mlx_norm_key], converted[mlx_norm_key])
            )

    def test_gemma4_convert_then_load_keeps_language_model_prefix(self):
        from mlx_lm.models import gemma4

        args = gemma4.ModelArgs.from_dict(
            {
                "model_type": "gemma4",
                "vocab_size": 32,
                "text_config": {
                    "model_type": "gemma4_text",
                    "hidden_size": 8,
                    "num_hidden_layers": 1,
                    "intermediate_size": 16,
                    "num_attention_heads": 1,
                    "num_key_value_heads": 1,
                    "num_global_key_value_heads": 1,
                    "head_dim": 8,
                    "global_head_dim": 8,
                    "sliding_window": 8,
                    "sliding_window_pattern": 1,
                    "layer_types": ["full_attention"],
                    "hidden_size_per_layer_input": 0,
                    "num_kv_shared_layers": 0,
                    "tie_word_embeddings": True,
                },
            }
        )
        model = gemma4.Model(args)

        base = mx.arange(8, dtype=mx.float32)
        hf_norm_key = "model.language_model.layers.0.input_layernorm.weight"
        mlx_norm_key = "language_model.model.layers.0.input_layernorm.weight"

        converted = model.sanitize(
            {
                hf_norm_key: base,
                "model.vision_tower.stub": mx.zeros((1,), dtype=mx.float32),
            }
        )
        self.assertIn(mlx_norm_key, converted)
        self.assertNotIn(
            "language_model.model.model.layers.0.input_layernorm.weight", converted
        )
        self.assertTrue(mx.array_equal(converted[mlx_norm_key], base))
        self.assertFalse(any("vision_tower" in k for k in converted))

        loaded = model.sanitize({mlx_norm_key: base})
        self.assertIn(mlx_norm_key, loaded)
        self.assertNotIn(
            "language_model.model.model.layers.0.input_layernorm.weight", loaded
        )
        self.assertTrue(mx.array_equal(loaded[mlx_norm_key], base))

    def test_gemma4_raw_hf_language_model_prefixes_model(self):
        from mlx_lm.models import gemma4

        args = gemma4.ModelArgs.from_dict(
            {
                "model_type": "gemma4",
                "vocab_size": 32,
                "text_config": {
                    "model_type": "gemma4_text",
                    "hidden_size": 8,
                    "num_hidden_layers": 1,
                    "intermediate_size": 16,
                    "num_attention_heads": 1,
                    "num_key_value_heads": 1,
                    "num_global_key_value_heads": 1,
                    "head_dim": 8,
                    "global_head_dim": 8,
                    "sliding_window": 8,
                    "sliding_window_pattern": 1,
                    "layer_types": ["full_attention"],
                    "hidden_size_per_layer_input": 0,
                    "num_kv_shared_layers": 0,
                    "tie_word_embeddings": True,
                },
            }
        )
        model = gemma4.Model(args)

        base = mx.arange(8, dtype=mx.float32)
        hf_norm_key = "model.language_model.layers.0.input_layernorm.weight"
        mlx_norm_key = "language_model.model.layers.0.input_layernorm.weight"

        converted = model.sanitize({hf_norm_key: base})
        self.assertIn(mlx_norm_key, converted)
        self.assertTrue(mx.array_equal(converted[mlx_norm_key], base))

    def test_gemma4_raw_hf_moe_expert_weights_split_for_switch_glu(self):
        from mlx_lm.models import gemma4

        args = gemma4.ModelArgs.from_dict(
            {
                "model_type": "gemma4",
                "vocab_size": 32,
                "text_config": {
                    "model_type": "gemma4_text",
                    "hidden_size": 8,
                    "num_hidden_layers": 1,
                    "intermediate_size": 16,
                    "num_attention_heads": 1,
                    "num_key_value_heads": 1,
                    "num_global_key_value_heads": 1,
                    "head_dim": 8,
                    "global_head_dim": 8,
                    "sliding_window": 8,
                    "sliding_window_pattern": 1,
                    "layer_types": ["full_attention"],
                    "hidden_size_per_layer_input": 0,
                    "num_kv_shared_layers": 0,
                    "tie_word_embeddings": True,
                    "enable_moe_block": True,
                    "num_experts": 2,
                    "top_k_experts": 1,
                    "moe_intermediate_size": 3,
                },
            }
        )
        model = gemma4.Model(args)

        gate_up = mx.arange(2 * 6 * 8, dtype=mx.float32).reshape(2, 6, 8)
        down = mx.arange(2 * 8 * 3, dtype=mx.float32).reshape(2, 8, 3)

        converted = model.sanitize(
            {
                "model.language_model.layers.0.experts.gate_up_proj": gate_up,
                "model.language_model.layers.0.experts.down_proj": down,
            }
        )

        gate_key = "language_model.model.layers.0.experts.switch_glu.gate_proj.weight"
        up_key = "language_model.model.layers.0.experts.switch_glu.up_proj.weight"
        down_key = "language_model.model.layers.0.experts.switch_glu.down_proj.weight"

        self.assertIn(gate_key, converted)
        self.assertIn(up_key, converted)
        self.assertIn(down_key, converted)
        self.assertTrue(mx.array_equal(converted[gate_key], gate_up[:, :3, :]))
        self.assertTrue(mx.array_equal(converted[up_key], gate_up[:, 3:, :]))
        self.assertTrue(mx.array_equal(converted[down_key], down))
        self.assertFalse(any("gate_up_proj" in k for k in converted))

    def test_gemma4_moe_router_quantizes_to_8bit(self):
        from mlx_lm.models import gemma4
        from mlx_lm.models.switch_layers import QuantizedSwitchLinear
        from mlx_lm.utils import quantize_model

        args = gemma4.ModelArgs.from_dict(
            {
                "model_type": "gemma4",
                "vocab_size": 64,
                "text_config": {
                    "model_type": "gemma4_text",
                    "hidden_size": 64,
                    "num_hidden_layers": 1,
                    "intermediate_size": 128,
                    "moe_intermediate_size": 128,
                    "num_attention_heads": 1,
                    "num_key_value_heads": 1,
                    "num_global_key_value_heads": 1,
                    "head_dim": 64,
                    "global_head_dim": 64,
                    "sliding_window": 8,
                    "sliding_window_pattern": 1,
                    "layer_types": ["full_attention"],
                    "hidden_size_per_layer_input": 0,
                    "num_kv_shared_layers": 0,
                    "tie_word_embeddings": True,
                    "enable_moe_block": True,
                    "num_experts": 8,
                    "top_k_experts": 2,
                },
            }
        )
        model = gemma4.Model(args)
        model, config = quantize_model(
            model,
            {"model_type": "gemma4", "text_config": copy.deepcopy(args.text_config)},
            group_size=64,
            bits=4,
        )

        layer = model.language_model.model.layers[0]
        self.assertIsInstance(layer.router.proj, nn.QuantizedLinear)
        self.assertEqual(layer.router.proj.bits, 8)
        self.assertIsInstance(layer.experts.switch_glu.gate_proj, QuantizedSwitchLinear)
        self.assertEqual(layer.experts.switch_glu.gate_proj.bits, 4)
        self.assertEqual(
            config["quantization"]["language_model.model.layers.0.router.proj"]["bits"],
            8,
        )
        self.assertEqual(config["quantization"]["bits"], 4)

    def test_qwen2_moe(self):
        from mlx_lm.models import qwen2_moe

        args = qwen2_moe.ModelArgs(
            model_type="qwen2_moe",
            hidden_size=1024,
            num_hidden_layers=4,
            intermediate_size=2048,
            num_attention_heads=4,
            rms_norm_eps=1e-5,
            vocab_size=10_000,
            num_experts_per_tok=4,
            num_experts=16,
            moe_intermediate_size=1024,
            shared_expert_intermediate_size=2048,
        )
        model = qwen2_moe.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_qwen2(self):
        from mlx_lm.models import qwen2

        args = qwen2.ModelArgs(
            model_type="qwen2",
            hidden_size=1024,
            num_hidden_layers=4,
            intermediate_size=2048,
            num_attention_heads=4,
            num_key_value_heads=4,
            rms_norm_eps=1e-5,
            vocab_size=10_000,
        )
        model = qwen2.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_qwen(self):
        from mlx_lm.models import qwen

        args = qwen.ModelArgs(
            model_type="qwen",
        )
        model = qwen.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_plamo(self):
        from mlx_lm.models import plamo

        args = plamo.ModelArgs(
            model_type="plamo",
            hidden_size=1024,
            num_hidden_layers=4,
            intermediate_size=2048,
            num_attention_heads=8,
            rms_norm_eps=1e-5,
            vocab_size=10_000,
        )
        model = plamo.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_plamo2(self):
        from mlx_lm.models import plamo2

        args = plamo2.ModelArgs(
            model_type="plamo2",
            hidden_size=1024,
            num_hidden_layers=4,
            intermediate_size=2048,
            num_attention_heads=8,
            rms_norm_eps=1e-5,
            vocab_size=10_000,
        )
        model = plamo2.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_stablelm(self):
        from mlx_lm.models import stablelm

        args = stablelm.ModelArgs(
            model_type="stablelm",
            vocab_size=10_000,
            hidden_size=1024,
            num_attention_heads=4,
            num_hidden_layers=4,
            num_key_value_heads=2,
            partial_rotary_factor=1.0,
            intermediate_size=2048,
            layer_norm_eps=1e-2,
            rope_theta=10_000,
            use_qkv_bias=False,
        )
        model = stablelm.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

        # StableLM 2
        args = stablelm.ModelArgs(
            model_type="stablelm",
            vocab_size=10000,
            hidden_size=512,
            num_attention_heads=8,
            num_hidden_layers=4,
            num_key_value_heads=2,
            partial_rotary_factor=0.25,
            intermediate_size=1024,
            layer_norm_eps=1e-5,
            rope_theta=10000,
            use_qkv_bias=True,
        )
        model = stablelm.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_starcoder2(self):
        from mlx_lm.models import starcoder2

        args = starcoder2.ModelArgs(
            model_type="starcoder2",
            hidden_size=1024,
            num_hidden_layers=4,
            intermediate_size=2048,
            num_attention_heads=4,
            num_key_value_heads=4,
        )
        model = starcoder2.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_step3p5(self):
        from mlx_lm.models import step3p5

        args = step3p5.ModelArgs(
            model_type="step3p5",
            hidden_size=256,
            num_hidden_layers=4,
            vocab_size=1024,
            num_attention_heads=4,
            num_attention_groups=2,
            head_dim=64,
            intermediate_size=512,
            rms_norm_eps=1e-5,
            rope_theta=[10000.0, 10000.0, 10000.0, 10000.0],
            sliding_window=64,
            layer_types=[
                "full_attention",
                "sliding_attention",
                "sliding_attention",
                "full_attention",
            ],
            partial_rotary_factors=[0.5, 1.0, 1.0, 0.5],
            attention_other_setting={
                "num_attention_heads": 8,
                "num_attention_groups": 2,
            },
            use_head_wise_attn_gate=True,
            moe_num_experts=4,
            moe_top_k=2,
            moe_intermediate_size=256,
            share_expert_dim=256,
            moe_layers_enum="1,2,3",
        )
        model = step3p5.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_step3p5_make_cache_uses_rotating_for_sliding_layers(self):
        from mlx_lm.models import step3p5

        args = step3p5.ModelArgs(
            model_type="step3p5",
            hidden_size=256,
            num_hidden_layers=4,
            vocab_size=1024,
            num_attention_heads=4,
            num_attention_groups=2,
            head_dim=64,
            intermediate_size=512,
            rms_norm_eps=1e-5,
            rope_theta=[10000.0, 10000.0, 10000.0, 10000.0],
            sliding_window=4,
            layer_types=[
                "full_attention",
                "sliding_attention",
                "sliding_attention",
                "full_attention",
            ],
            partial_rotary_factors=[0.5, 1.0, 1.0, 0.5],
            attention_other_setting={
                "num_attention_heads": 8,
                "num_attention_groups": 2,
            },
            use_head_wise_attn_gate=True,
            moe_num_experts=4,
            moe_top_k=2,
            moe_intermediate_size=256,
            share_expert_dim=256,
            moe_layers_enum="1,2,3",
        )
        model = step3p5.Model(args)

        caches = model.make_cache()
        self.assertIsInstance(caches[0], KVCache)
        self.assertIsInstance(caches[1], RotatingKVCache)
        self.assertIsInstance(caches[2], RotatingKVCache)
        self.assertIsInstance(caches[3], KVCache)

        tokens = mx.array([[1, 2, 3, 4, 5, 6, 7]], dtype=mx.int32)
        step = model(tokens[:, :3], cache=caches)
        mx.eval(step)
        for i in range(3, 7):
            step = model(tokens[:, i : i + 1], cache=caches)
            mx.eval(step)

        self.assertEqual(caches[0].size(), 7)
        self.assertEqual(caches[1].size(), args.sliding_window)
        self.assertEqual(caches[2].size(), args.sliding_window)
        self.assertEqual(caches[3].size(), 7)

    def test_step3p5_make_cache_uses_fallback_sliding_pattern(self):
        from mlx_lm.models import step3p5

        args = step3p5.ModelArgs(
            model_type="step3p5",
            hidden_size=256,
            num_hidden_layers=5,
            vocab_size=1024,
            num_attention_heads=4,
            num_attention_groups=2,
            head_dim=64,
            intermediate_size=512,
            rms_norm_eps=1e-5,
            rope_theta=10000.0,
            sliding_window=4,
            partial_rotary_factors=[1.0] * 5,
            use_head_wise_attn_gate=True,
            moe_num_experts=4,
            moe_top_k=2,
            moe_intermediate_size=256,
            share_expert_dim=256,
            moe_layers_enum="1,2,3,4",
        )
        model = step3p5.Model(args)

        caches = model.make_cache()
        self.assertIsInstance(caches[0], RotatingKVCache)
        self.assertIsInstance(caches[1], KVCache)
        self.assertIsInstance(caches[2], RotatingKVCache)
        self.assertIsInstance(caches[3], KVCache)
        self.assertIsInstance(caches[4], RotatingKVCache)

        tokens = mx.array([[1, 2, 3, 4, 5, 6]], dtype=mx.int32)
        step = model(tokens[:, :2], cache=caches)
        mx.eval(step)
        for i in range(2, 6):
            step = model(tokens[:, i : i + 1], cache=caches)
            mx.eval(step)

        self.assertEqual(caches[0].size(), args.sliding_window)
        self.assertEqual(caches[1].size(), 6)
        self.assertEqual(caches[2].size(), args.sliding_window)
        self.assertEqual(caches[3].size(), 6)
        self.assertEqual(caches[4].size(), args.sliding_window)

    def test_cohere(self):
        from mlx_lm.models import cohere

        args = cohere.ModelArgs(
            model_type="cohere",
        )
        model = cohere.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_dbrx(self):
        from mlx_lm.models import dbrx

        args = dbrx.ModelArgs(
            model_type="dbrx",
            d_model=1024,
            ffn_config={"ffn_hidden_size": 2048, "moe_num_experts": 4, "moe_top_k": 2},
            attn_config={"kv_n_heads": 2, "clip_qkv": True, "rope_theta": 10000},
            n_layers=4,
            n_heads=4,
            vocab_size=10_000,
        )
        model = dbrx.Model(args)
        self.model_test_runner(model, args.model_type, args.vocab_size, args.n_layers)

    def test_minicpm(self):
        from mlx_lm.models import minicpm

        args = minicpm.ModelArgs(
            model_type="minicpm",
            hidden_size=1024,
            dim_model_base=1024,
            num_hidden_layers=4,
            intermediate_size=2048,
            num_attention_heads=4,
            rms_norm_eps=1e-4,
            vocab_size=10000,
            num_key_value_heads=2,
            scale_depth=1.0,
            scale_emb=1.0,
        )
        model = minicpm.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_mamba(self):
        from mlx_lm.models import mamba

        args = mamba.ModelArgs(
            model_type="mamba",
            vocab_size=10000,
            use_bias=False,
            use_conv_bias=True,
            conv_kernel=4,
            hidden_size=768,
            num_hidden_layers=24,
            state_size=16,
            intermediate_size=1536,
            time_step_rank=48,
        )
        model = mamba.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_falcon_h1(self):
        from mlx_lm.models import falcon_h1

        args = falcon_h1.ModelArgs(
            model_type="falcon_h1",
            num_hidden_layers=12,
            vocab_size=10000,
        )
        model = falcon_h1.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_gpt2(self):
        from mlx_lm.models import gpt2

        args = gpt2.ModelArgs(
            model_type="gpt2",
            n_ctx=1024,
            n_embd=768,
            n_head=12,
            n_layer=12,
            n_positions=1024,
            layer_norm_epsilon=1e-5,
            vocab_size=50256,
        )
        model = gpt2.Model(args)
        self.model_test_runner(model, args.model_type, args.vocab_size, args.n_layer)

    def test_gpt_neox(self):
        from mlx_lm.models import gpt_neox

        args = gpt_neox.ModelArgs(
            model_type="gpt_neox",
            max_position_embeddings=2048,
            hidden_size=6144,
            num_attention_heads=64,
            num_hidden_layers=44,
            layer_norm_eps=1e-5,
            vocab_size=50432,
            rotary_emb_base=10_000,
            rotary_pct=0.25,
        )
        model = gpt_neox.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_openelm(self):
        from mlx_lm.models import openelm

        args = openelm.ModelArgs(
            model_type="openelm",
            ffn_dim_divisor=256,
            ffn_multipliers=[
                0.5,
                0.73,
                0.97,
                1.2,
                1.43,
                1.67,
                1.9,
                2.13,
                2.37,
                2.6,
                2.83,
                3.07,
                3.3,
                3.53,
                3.77,
                4.0,
            ],
            head_dim=64,
            model_dim=1280,
            normalize_qk_projections=True,
            num_kv_heads=[3, 3, 3, 3, 3, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5],
            num_query_heads=[
                12,
                12,
                12,
                12,
                12,
                16,
                16,
                16,
                16,
                16,
                16,
                16,
                20,
                20,
                20,
                20,
            ],
            num_transformer_layers=16,
            vocab_size=32000,
        )

        model = openelm.Model(args)
        self.model_test_runner(
            model,
            args.model_type,
            args.vocab_size,
            len(args.ffn_multipliers),
        )

    def test_internlm2(self):
        from mlx_lm.models import internlm2

        args = internlm2.ModelArgs(
            model_type="internlm2",
            hidden_size=1024,
            num_hidden_layers=4,
            intermediate_size=2048,
            num_attention_heads=4,
            rms_norm_eps=1e-5,
            vocab_size=10000,
        )
        model = internlm2.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_llama3_1(self):
        from mlx_lm.models import llama

        args = llama.ModelArgs(
            model_type="llama",
            hidden_size=1024,
            num_hidden_layers=4,
            intermediate_size=2048,
            num_attention_heads=4,
            rms_norm_eps=1e-5,
            vocab_size=10_000,
            max_position_embeddings=128,
            mlp_bias=False,
            num_key_value_heads=2,
            rope_scaling={
                "factor": 8.0,
                "low_freq_factor": 1.0,
                "high_freq_factor": 4.0,
                "original_max_position_embeddings": 8192,
                "rope_type": "llama3",
            },
        )
        model = llama.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_deepseek(self):
        from mlx_lm.models import deepseek

        args = deepseek.ModelArgs(
            model_type="deepseek",
            vocab_size=1024,
            hidden_size=128,
            intermediate_size=256,
            moe_intermediate_size=256,
            num_hidden_layers=4,
            num_attention_heads=8,
            num_key_value_heads=4,
        )
        model = deepseek.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_deepseek_v2(self):
        from mlx_lm.models import deepseek_v2

        args = deepseek_v2.ModelArgs(
            model_type="deepseek_v2",
            vocab_size=1024,
            hidden_size=128,
            intermediate_size=256,
            moe_intermediate_size=256,
            num_hidden_layers=4,
            num_attention_heads=4,
            num_key_value_heads=2,
            kv_lora_rank=4,
            q_lora_rank=4,
            qk_rope_head_dim=32,
            v_head_dim=16,
            qk_nope_head_dim=32,
            rope_scaling={
                "beta_fast": 32,
                "beta_slow": 1,
                "factor": 40,
                "mscale": 1.0,
                "mscale_all_dim": 1.0,
                "original_max_position_embeddings": 4096,
                "type": "yarn",
            },
        )
        model = deepseek_v2.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_deepseek_v3(self):
        from mlx_lm.models import deepseek_v3

        args = deepseek_v3.ModelArgs(
            model_type="deepseek_v3",
            vocab_size=1024,
            hidden_size=128,
            intermediate_size=256,
            moe_intermediate_size=256,
            num_hidden_layers=4,
            num_attention_heads=4,
            num_key_value_heads=2,
            n_routed_experts=4,
            n_group=2,
            topk_group=1,
            num_experts_per_tok=2,
            n_shared_experts=1,
            kv_lora_rank=4,
            q_lora_rank=4,
            qk_rope_head_dim=32,
            v_head_dim=16,
            qk_nope_head_dim=32,
            rope_scaling={
                "beta_fast": 32,
                "beta_slow": 1,
                "factor": 40,
                "mscale": 1.0,
                "mscale_all_dim": 1.0,
                "original_max_position_embeddings": 4096,
                "type": "yarn",
            },
        )
        model = deepseek_v3.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_gemma2(self):
        from mlx_lm.models import gemma2

        args = gemma2.ModelArgs(
            model_type="gemma2",
            hidden_size=128,
            num_hidden_layers=4,
            intermediate_size=256,
            num_attention_heads=2,
            head_dim=32,
            rms_norm_eps=1e-4,
            vocab_size=1024,
            num_key_value_heads=2,
        )
        model = gemma2.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_gemma3_text(self):
        from mlx_lm.models import gemma3_text

        args = gemma3_text.ModelArgs(
            model_type="gemma3_text",
            hidden_size=128,
            num_hidden_layers=12,
            intermediate_size=256,
            num_attention_heads=4,
            head_dim=32,
            rms_norm_eps=1e-4,
            num_key_value_heads=1,
            sliding_window=1024,
            sliding_window_pattern=6,
        )
        model = gemma3_text.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_gemma4_text(self):
        from mlx_lm.models import gemma4_text

        args = gemma4_text.ModelArgs(
            model_type="gemma4_text",
            hidden_size=128,
            num_hidden_layers=10,
            intermediate_size=256,
            num_attention_heads=4,
            head_dim=32,
            global_head_dim=64,
            rms_norm_eps=1e-6,
            vocab_size=1000,
            vocab_size_per_layer_input=1000,
            num_key_value_heads=1,
            num_kv_shared_layers=4,
            hidden_size_per_layer_input=32,
            sliding_window=8,
            sliding_window_pattern=5,
            final_logit_softcapping=30.0,
            layer_types=[
                "sliding_attention",
                "sliding_attention",
                "sliding_attention",
                "sliding_attention",
                "full_attention",
                "sliding_attention",
                "sliding_attention",
                "sliding_attention",
                "sliding_attention",
                "full_attention",
            ],
            rope_parameters={
                "full_attention": {
                    "partial_rotary_factor": 0.25,
                    "rope_theta": 1000000.0,
                },
                "sliding_attention": {
                    "rope_theta": 10000.0,
                },
            },
        )
        model = gemma4_text.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_gemma4_quantized_embedding_preserves_lookup_scale(self):
        from mlx_lm.models import gemma4_text

        args = gemma4_text.ModelArgs(
            model_type="gemma4_text",
            hidden_size=32,
            num_hidden_layers=1,
            intermediate_size=64,
            num_attention_heads=2,
            num_key_value_heads=1,
            num_global_key_value_heads=1,
            head_dim=16,
            global_head_dim=16,
            sliding_window=8,
            sliding_window_pattern=1,
            layer_types=["full_attention"],
            hidden_size_per_layer_input=0,
            vocab_size=4,
            num_kv_shared_layers=0,
        )
        model = gemma4_text.Gemma4TextModel(args)
        model.embed_tokens.weight = mx.ones((4, 32), dtype=mx.float32)
        model.embed_tokens = model.embed_tokens.to_quantized(group_size=32, bits=8)

        token_ids = mx.array([[0, 1]], dtype=mx.int32)
        lookup = model.embed_tokens(token_ids) * model.embed_scale
        logits = model.embed_tokens.as_linear(mx.ones((1, 1, 32), dtype=mx.float32))
        mx.eval(lookup, logits)

        self.assertTrue(
            mx.allclose(
                lookup,
                mx.ones((1, 2, 32), dtype=mx.float32) * (32.0**0.5),
            )
        )
        self.assertTrue(
            mx.allclose(logits, mx.ones((1, 1, 4), dtype=mx.float32) * 32.0)
        )

    def test_gemma4_kv_shared_layers_omit_kv_projections(self):
        """KV-shared layers must not create k_proj/v_proj/k_norm/v_norm so that
        models saved without redundant weights (e.g. via transformers
        save_pretrained) can be loaded with strict=True."""
        from mlx_lm.models import gemma4_text

        args = gemma4_text.ModelArgs(
            model_type="gemma4_text",
            hidden_size=128,
            num_hidden_layers=10,
            intermediate_size=256,
            num_attention_heads=4,
            head_dim=32,
            global_head_dim=64,
            rms_norm_eps=1e-6,
            vocab_size=1000,
            vocab_size_per_layer_input=1000,
            num_key_value_heads=1,
            num_kv_shared_layers=4,
            hidden_size_per_layer_input=32,
            sliding_window=8,
            sliding_window_pattern=5,
            final_logit_softcapping=30.0,
            layer_types=[
                "sliding_attention",
                "sliding_attention",
                "sliding_attention",
                "sliding_attention",
                "full_attention",
                "sliding_attention",
                "sliding_attention",
                "sliding_attention",
                "sliding_attention",
                "full_attention",
            ],
            rope_parameters={
                "full_attention": {
                    "partial_rotary_factor": 0.25,
                    "rope_theta": 1000000.0,
                },
                "sliding_attention": {
                    "rope_theta": 10000.0,
                },
            },
        )
        model = gemma4_text.Model(args)

        # Non-shared layers (0-5) should have KV projections
        for i in range(6):
            attn = model.model.layers[i].self_attn
            self.assertTrue(attn.has_kv)
            self.assertTrue(hasattr(attn, "k_proj"))
            self.assertTrue(hasattr(attn, "k_norm"))

        # Shared layers (6-9) should NOT have KV projections
        for i in range(6, 10):
            attn = model.model.layers[i].self_attn
            self.assertFalse(attn.has_kv)
            self.assertFalse(hasattr(attn, "k_proj"))
            self.assertFalse(hasattr(attn, "k_norm"))
            self.assertFalse(hasattr(attn, "v_proj"))

        # Verify the model can load weights that omit shared-layer KV params
        weights = dict(tree_flatten(model.parameters()))
        kv_keys = [
            k for k in weights if "k_proj" in k or "v_proj" in k or "k_norm" in k
        ]
        for k in kv_keys:
            # All KV keys should belong to non-shared layers (0-5)
            layer_idx = int(k.split("layers.")[1].split(".")[0])
            self.assertLess(layer_idx, 6)

    def test_gemma4_input_embeddings_reconstruct_per_layer_inputs(self):
        from mlx_lm.models import gemma4_text

        args = gemma4_text.ModelArgs(
            model_type="gemma4_text",
            hidden_size=32,
            num_hidden_layers=2,
            intermediate_size=64,
            num_attention_heads=2,
            num_key_value_heads=1,
            num_global_key_value_heads=1,
            head_dim=16,
            global_head_dim=16,
            sliding_window=8,
            sliding_window_pattern=1,
            layer_types=["full_attention", "full_attention"],
            hidden_size_per_layer_input=8,
            vocab_size=32,
            vocab_size_per_layer_input=32,
            num_kv_shared_layers=0,
        )
        model = gemma4_text.Model(args)
        tokens = mx.array([[1, 2, 3]], dtype=mx.int32)
        embeddings = model.model.embed_tokens(tokens)
        per_layer_inputs = model.model._get_per_layer_inputs(tokens)

        direct = model(tokens)
        from_embeddings = model(None, input_embeddings=embeddings)
        explicit = model(
            None,
            input_embeddings=embeddings,
            per_layer_inputs=per_layer_inputs,
        )
        mx.eval(direct, from_embeddings, explicit)

        self.assertTrue(
            mx.allclose(direct.astype(mx.float32), from_embeddings.astype(mx.float32))
        )
        self.assertTrue(
            mx.allclose(direct.astype(mx.float32), explicit.astype(mx.float32))
        )

    def test_gpt_bigcode(self):
        from mlx_lm.models import gpt_bigcode

        args = gpt_bigcode.ModelArgs(
            model_type="gpt_bigcode",
            n_embd=128,
            n_layer=128,
            n_inner=256,
            n_head=4,
            n_positions=1000,
            layer_norm_epsilon=1e-5,
            vocab_size=1024,
        )
        model = gpt_bigcode.Model(args)
        self.model_test_runner(model, args.model_type, args.vocab_size, args.n_layer)

    def test_nemotron(self):
        from mlx_lm.models import nemotron

        args = nemotron.ModelArgs(
            model_type="nemotron",
            hidden_size=128,
            hidden_act="gelu",
            num_hidden_layers=4,
            intermediate_size=256,
            num_attention_heads=4,
            norm_eps=1e-5,
            vocab_size=1024,
            num_key_value_heads=2,
        )
        model = nemotron.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_phi3small(self):
        from mlx_lm.models import phi3small

        args = phi3small.ModelArgs(
            model_type="phi3small",
            hidden_size=128,
            dense_attention_every_n_layers=2,
            ff_intermediate_size=256,
            gegelu_limit=1.0,
            num_hidden_layers=4,
            num_attention_heads=4,
            num_key_value_heads=2,
            layer_norm_epsilon=1e-4,
            vocab_size=1000,
        )
        model = phi3small.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_phimoe(self):
        from mlx_lm.models import phimoe

        args = phimoe.ModelArgs(
            model_type="phimoe",
            vocab_size=320,
            hidden_size=128,
            intermediate_size=256,
            num_hidden_layers=4,
            num_attention_heads=4,
            num_key_value_heads=4,
            rope_scaling={
                "long_factor": [1.0] * 16,
                "long_mscale": 1.243163121016122,
                "original_max_position_embeddings": 4096,
                "short_factor": [1.0] * 16,
                "short_mscale": 1.243163121016122,
                "type": "longrope",
            },
        )
        model = phimoe.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_recurrent_gemma(self):
        from mlx_lm.models import recurrent_gemma

        args = recurrent_gemma.ModelArgs(
            model_type="recurrent_gemma",
            hidden_size=128,
            attention_bias=False,
            conv1d_width=3,
            intermediate_size=256,
            logits_soft_cap=1.0,
            num_attention_heads=4,
            num_hidden_layers=4,
            num_key_value_heads=2,
            rms_norm_eps=1e-4,
            rope_theta=1000,
            attention_window_size=1024,
            vocab_size=1000,
            block_types=["recurrent", "recurrent", "attention"],
        )
        model = recurrent_gemma.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_hunyuan(self):
        from mlx_lm.models import hunyuan

        args = hunyuan.ModelArgs(
            model_type="hunyuan",
            hidden_size=128,
            attention_bias=False,
            intermediate_size=256,
            num_attention_heads=4,
            num_hidden_layers=4,
            num_key_value_heads=2,
            rms_norm_eps=1e-4,
            rope_theta=1000,
            vocab_size=1000,
            moe_topk=2,
            num_experts=2,
            num_shared_expert=1,
            use_mixed_mlp_moe=True,
            use_qk_norm=True,
            rope_scaling={
                "alpha": 1000.0,
                "factor": 1.0,
                "type": "dynamic",
            },
            use_cla=True,
            cla_share_factor=2,
        )
        model = hunyuan.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_hunyuan_v1_dense(self):
        from mlx_lm.models import hunyuan_v1_dense

        args = hunyuan_v1_dense.ModelArgs(
            model_type="hunyuan_v1_dense",
            hidden_size=128,
            attention_bias=False,
            intermediate_size=256,
            num_attention_heads=4,
            num_hidden_layers=4,
            num_key_value_heads=2,
            rms_norm_eps=1e-4,
            rope_theta=1000,
            vocab_size=1000,
            use_qk_norm=True,
            rope_scaling={
                "alpha": 1000.0,
                "factor": 1.0,
                "type": "dynamic",
                "beta_fast": 32,
                "beta_slow": 1,
                "mscale": 1.0,
                "mscale_all_dim": 0.0,
                "original_max_position_embeddings": 8192,
            },
            max_position_embeddings=32768,
        )
        model = hunyuan_v1_dense.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_olmo2(self):
        from mlx_lm.models import olmo2

        args = olmo2.ModelArgs(
            model_type="olmo2",
            hidden_size=128,
            attention_bias=False,
            intermediate_size=256,
            num_attention_heads=4,
            num_hidden_layers=4,
            num_key_value_heads=2,
            rms_norm_eps=1e-4,
            rope_theta=1000,
            vocab_size=1000,
        )
        model = olmo2.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_exaone(self):
        from mlx_lm.models import exaone

        args = exaone.ModelArgs(
            model_type="exaone",
            hidden_size=128,
            num_layers=4,
            intermediate_size=256,
            num_attention_heads=8,
            num_key_value_heads=2,
            vocab_size=1000,
            layer_norm_epsilon=1e-4,
            rope_theta=10000,
        )
        model = exaone.Model(args)
        self.model_test_runner(model, args.model_type, args.vocab_size, args.num_layers)

    def test_cohere2(self):
        from mlx_lm.models import cohere2

        args = cohere2.ModelArgs(
            model_type="cohere2",
            hidden_size=4096,
            head_dim=128,
            num_hidden_layers=40,
            sliding_window=4096,
            sliding_window_pattern=4,
        )
        model = cohere2.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_internlm3(self):
        from mlx_lm.models import internlm3

        args = internlm3.ModelArgs(
            model_type="internlm3",
            hidden_size=1024,
            num_hidden_layers=4,
            intermediate_size=2048,
            num_attention_heads=4,
            rms_norm_eps=1e-5,
            vocab_size=10_000,
        )
        model = internlm3.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_iquestloopcoder(self):
        from mlx_lm.models import iquestloopcoder

        args = iquestloopcoder.ModelArgs(
            model_type="iquestloopcoder",
            hidden_size=256,
            num_hidden_layers=2,
            intermediate_size=512,
            num_attention_heads=4,
            num_key_value_heads=2,
            rms_norm_eps=1e-5,
            head_dim=32,
            vocab_size=1000,
            rope_theta=500000.0,
            tie_word_embeddings=False,
            loop_num=2,
            loop_window_size=32,
        )
        model = iquestloopcoder.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_smollm3(self):
        from mlx_lm.models import smollm3

        args = smollm3.ModelArgs(
            model_type="smollm3",
            hidden_size=1024,
            num_hidden_layers=4,
            intermediate_size=2048,
            num_attention_heads=4,
            rms_norm_eps=1e-5,
            vocab_size=10_000,
        )
        model = smollm3.Model(args)
        self.model_test_runner(
            model, "smollm3", args.vocab_size, args.num_hidden_layers
        )

    def test_gpt_oss(self):
        from mlx_lm.models import gpt_oss

        args = gpt_oss.ModelArgs(
            model_type="gpt_oss",
            hidden_size=1024,
            num_hidden_layers=4,
            intermediate_size=2048,
            num_attention_heads=8,
            num_key_value_heads=2,
            num_local_experts=16,
            num_experts_per_tok=2,
            sliding_window=128,
            rope_theta=10000,
            vocab_size=10_000,
            layer_types=[
                "sliding_attention",
                "full_attention",
                "sliding_attention",
                "full_attention",
            ],
        )
        model = gpt_oss.Model(args)
        self.model_test_runner(
            model, args.model_type, args.vocab_size, args.num_hidden_layers
        )

    def test_all_models(self):
        test_configs = [
            {
                "model_type": "afm7",
                "vocab_size": 1000,
                "hidden_dim": 256,
                "num_layers": 16,
                "num_hidden_layers": 16,
                "num_kv_reuse_layers": 8,
                "num_heads": 8,
                "num_kv_heads": 4,
            },
            {
                "model_type": "apertus",
                "hidden_size": 128,
                "num_hidden_layers": 4,
                "intermediate_size": 256,
                "mlp_bias": True,
                "num_attention_heads": 8,
                "attention_bias": False,
                "rms_norm_eps": 1e-5,
                "vocab_size": 1000,
                "num_key_value_heads": 4,
                "max_position_embeddings": 1000,
                "rope_theta": 1000,
                "post_norm": True,
                "qk_norm": True,
                "tie_word_embeddings": False,
            },
            {
                "model_type": "baichuan_m1",
                "vocab_size": 1000,
                "hidden_size": 128,
                "intermediate_size": 256,
                "num_hidden_layers": 4,
                "num_attention_heads": 4,
                "num_key_value_heads": 2,
                "rope_theta": 1000,
                "sliding_window": 128,
                "sliding_window_layers": [0, 2],
                "conv_window": 2,
                "rms_norm_eps": 1e-5,
            },
            {
                "model_type": "bailing_moe",
                "hidden_size": 64,
                "intermediate_size": 128,
                "max_position_embeddings": 1000,
                "moe_intermediate_size": 128,
                "num_experts": 4,
                "num_shared_experts": 1,
                "norm_topk_prob": True,
                "num_attention_heads": 4,
                "num_experts_per_tok": 2,
                "num_hidden_layers": 4,
                "num_key_value_heads": 2,
                "rms_norm_eps": 1e-5,
                "rope_theta": 1000,
                "vocab_size": 1000,
                "first_k_dense_replace": 2,
            },
            {
                "model_type": "dots1",
                "hidden_size": 64,
                "num_hidden_layers": 4,
                "intermediate_size": 128,
                "num_attention_heads": 4,
                "rms_norm_eps": 1e-5,
                "vocab_size": 1000,
                "max_position_embeddings": None,
                "num_key_value_heads": 2,
                "first_k_dense_replace": 1,
                "moe_intermediate_size": 64,
                "n_routed_experts": 4,
                "n_shared_experts": 1,
                "norm_topk_prob": True,
                "num_experts_per_tok": 1,
                "rope_theta": 1000,
                "routed_scaling_factor": 1.0,
            },
            {
                "hidden_size": 128,
                "intermediate_size": 128,
                "model_type": "ernie4_5",
                "max_position_embeddings": 1000,
                "num_attention_heads": 4,
                "num_key_value_heads": 2,
                "head_dim": None,
                "num_hidden_layers": 4,
                "rms_norm_eps": 1e-5,
                "vocab_size": 1000,
                "rope_theta": 10000,
                "use_bias": False,
                "tie_word_embeddings": True,
            },
            {
                "hidden_size": 128,
                "intermediate_size": 128,
                "model_type": "ernie4_5_moe",
                "max_position_embeddings": 1000,
                "num_attention_heads": 4,
                "num_key_value_heads": 2,
                "num_hidden_layers": 4,
                "rms_norm_eps": 1e-5,
                "vocab_size": 1000,
                "rope_theta": 1000,
                "use_bias": False,
                "tie_word_embeddings": False,
                "moe_num_experts": 4,
            },
            {
                "model_type": "exaone4",
                "hidden_size": 128,
                "intermediate_size": 128,
                "num_attention_heads": 4,
                "vocab_size": 1000,
                "rms_norm_eps": 1e-5,
                "num_hidden_layers": 4,
                "max_position_embeddings": 1000,
                "rope_theta": 10000,
                "layer_norm_epsilon": 1e-5,
                "num_key_value_heads": 2,
                "head_dim": 32,
                "tie_word_embeddings": False,
                "rope_scaling": None,
                "sliding_window": 8,
                "sliding_window_pattern": "LLGL",
            },
            {
                "model_type": "gemma4",
                "num_hidden_layers": 10,
                "vocab_size": 1000,
                "text_config": {
                    "model_type": "gemma4_text",
                    "hidden_size": 128,
                    "num_hidden_layers": 10,
                    "intermediate_size": 128,
                    "num_attention_heads": 4,
                    "head_dim": 32,
                    "global_head_dim": 64,
                    "rms_norm_eps": 1e-6,
                    "vocab_size": 1000,
                    "vocab_size_per_layer_input": 1000,
                    "num_key_value_heads": 1,
                    "num_kv_shared_layers": 4,
                    "hidden_size_per_layer_input": 32,
                    "sliding_window": 8,
                    "sliding_window_pattern": 5,
                    "final_logit_softcapping": 30.0,
                    "layer_types": [
                        "sliding_attention",
                        "sliding_attention",
                        "sliding_attention",
                        "sliding_attention",
                        "full_attention",
                        "sliding_attention",
                        "sliding_attention",
                        "sliding_attention",
                        "sliding_attention",
                        "full_attention",
                    ],
                    "rope_parameters": {
                        "full_attention": {
                            "partial_rotary_factor": 0.25,
                            "rope_theta": 1000000.0,
                        },
                        "sliding_attention": {
                            "rope_theta": 10000.0,
                        },
                    },
                },
            },
            {
                "model_type": "gemma3n",
                "num_hidden_layers": 4,
                "vocab_size": 1000,
                "text_config": {
                    "model_type": "gemma3n",
                    "hidden_size": 128,
                    "num_hidden_layers": 4,
                    "intermediate_size": 128,
                    "num_attention_heads": 4,
                    "head_dim": 32,
                    "rms_norm_eps": 1e-5,
                    "vocab_size": 1000,
                    "num_key_value_heads": 2,
                    "num_kv_shared_layers": 2,
                    "vocab_size_per_layer_input": 1000,
                    "sliding_window": 8,
                    "max_position_embeddings": 1000,
                    "rope_local_base_freq": 1.0,
                    "rope_theta": 1000.0,
                    "final_logit_softcapping": 1.0,
                    "layer_types": [
                        "sliding_attention",
                        "full_attention",
                        "sliding_attention",
                        "full_attention",
                    ],
                    "activation_sparsity_pattern": [0.5, 0.5, 0.5, 0.5],
                    "hidden_size_per_layer_input": 256,
                    "altup_num_inputs": 4,
                    "altup_coef_clip": 1.0,
                    "altup_correct_scale": True,
                    "altup_active_idx": 0,
                    "laurel_rank": 8,
                },
            },
            {
                "model_type": "glm4",
                "hidden_size": 256,
                "num_hidden_layers": 4,
                "intermediate_size": 128,
                "num_attention_heads": 4,
                "attention_bias": False,
                "head_dim": 64,
                "rms_norm_eps": 1e-5,
                "vocab_size": 1000,
                "num_key_value_heads": 2,
                "partial_rotary_factor": 0.5,
                "rope_theta": 1000,
            },
            {
                "model_type": "glm4_moe",
                "vocab_size": 1000,
                "hidden_size": 128,
                "intermediate_size": 128,
                "max_position_embeddings": 1000,
                "moe_intermediate_size": 128,
                "norm_topk_prob": True,
                "num_attention_heads": 4,
                "n_group": 2,
                "head_dim": 32,
                "topk_group": 1,
                "n_shared_experts": 1,
                "n_routed_experts": 4,
                "routed_scaling_factor": 1.0,
                "num_experts_per_tok": 2,
                "first_k_dense_replace": 1,
                "num_hidden_layers": 4,
                "num_key_value_heads": 2,
                "rms_norm_eps": 1e-5,
                "rope_theta": 1000,
                "rope_scaling": None,
                "use_qk_norm": True,
                "tie_word_embeddings": False,
                "attention_bias": False,
                "partial_rotary_factor": 0.5,
            },
            {
                "model_type": "glm4_moe_lite",
                "vocab_size": 1000,
                "hidden_size": 64,
                "intermediate_size": 128,
                "moe_intermediate_size": 32,
                "num_hidden_layers": 4,
                "num_attention_heads": 4,
                "num_key_value_heads": 4,
                "n_shared_experts": 1,
                "n_routed_experts": 4,
                "routed_scaling_factor": 1.0,
                "kv_lora_rank": 8,
                "q_lora_rank": 8,
                "qk_rope_head_dim": 8,
                "qk_nope_head_dim": 16,
                "v_head_dim": 8,
                "topk_method": "noaux_tc",
                "scoring_func": "sigmoid",
                "norm_topk_prob": True,
                "n_group": 1,
                "topk_group": 1,
                "num_experts_per_tok": 2,
                "moe_layer_freq": 1,
                "first_k_dense_replace": 1,
                "max_position_embeddings": 256,
                "rms_norm_eps": 1e-5,
                "rope_theta": 1000,
                "rope_scaling": None,
                "attention_bias": False,
                "partial_rotary_factor": 1.0,
                "tie_word_embeddings": False,
                "num_nextn_predict_layers": 1,
            },
            {
                "model_type": "granite",
                "hidden_size": 128,
                "num_hidden_layers": 4,
                "intermediate_size": 128,
                "num_attention_heads": 4,
                "rms_norm_eps": 1e-5,
                "vocab_size": 1000,
                "logits_scaling": 1.0,
                "attention_multiplier": 1.0,
                "embedding_multiplier": 1.0,
                "residual_multiplier": 1.0,
                "max_position_embeddings": 1000,
                "num_key_value_heads": 2,
                "attention_bias": False,
                "mlp_bias": False,
                "rope_theta": 1000,
            },
            {
                "model_type": "granitemoe",
                "hidden_size": 128,
                "num_hidden_layers": 4,
                "intermediate_size": 128,
                "num_attention_heads": 4,
                "rms_norm_eps": 1e-5,
                "vocab_size": 1000,
                "logits_scaling": 1.0,
                "attention_multiplier": 1.0,
                "embedding_multiplier": 1.0,
                "residual_multiplier": 1.0,
                "max_position_embeddings": 1000,
                "num_key_value_heads": 2,
                "attention_bias": False,
                "rope_theta": 1000,
                "num_local_experts": 4,
                "num_experts_per_tok": 2,
            },
            {
                "hidden_size": 256,
                "num_hidden_layers": 4,
                "intermediate_size": 256,
                "num_attention_heads": 4,
                "num_key_value_heads": 2,
                "rms_norm_eps": 1e-5,
                "vocab_size": 1000,
                "attention_bias": False,
                "head_dim": 64,
                "max_position_embeddings": 1000,
                "mlp_bias": False,
                "model_type": "helium",
                "rope_theta": 1000,
                "tie_word_embeddings": False,
            },
            {
                "text_config": {
                    "vocab_size": 1000,
                    "hidden_size": 128,
                    "intermediate_size": 128,
                    "moe_intermediate_size": 128,
                    "num_hidden_layers": 4,
                    "num_attention_heads": 4,
                    "num_key_value_heads": 2,
                    "n_shared_experts": 1,
                    "n_routed_experts": 4,
                    "kv_lora_rank": 32,
                    "q_lora_rank": 32,
                    "qk_rope_head_dim": 8,
                    "v_head_dim": 16,
                    "qk_nope_head_dim": 16,
                },
                "model_type": "kimi_vl",
                "num_hidden_layers": 4,
                "vocab_size": 1000,
            },
            {
                "model_type": "lfm2-vl",
                "vocab_size": 1000,
                "num_hidden_layers": 4,
                "text_config": {
                    "model_type": "lfm2",
                    "vocab_size": 1000,
                    "hidden_size": 128,
                    "num_hidden_layers": 4,
                    "num_attention_heads": 4,
                    "num_key_value_heads": 2,
                    "max_position_embeddings": 1000,
                    "norm_eps": 1e-5,
                    "conv_bias": False,
                    "conv_L_cache": 3,
                    "block_dim": 128,
                    "block_ff_dim": 128,
                    "block_multiple_of": 4,
                    "block_ffn_dim_multiplier": 2,
                    "block_auto_adjust_ff_dim": True,
                    "layer_types": ["full_attention", "", "full_attention", ""],
                    "rope_theta": 1000,
                },
            },
            {
                "model_type": "llama4",
                "text_config": {
                    "attention_bias": False,
                    "attention_chunk_size": 8,
                    "head_dim": 32,
                    "hidden_size": 128,
                    "interleave_moe_layer_step": 2,
                    "intermediate_size": 128,
                    "intermediate_size_mlp": 128,
                    "max_position_embeddings": 1000,
                    "model_type": "llama4",
                    "num_attention_heads": 4,
                    "num_experts_per_tok": 1,
                    "num_hidden_layers": 4,
                    "num_key_value_heads": 2,
                    "num_local_experts": 2,
                    "rms_norm_eps": 1e-4,
                    "rope_scaling": None,
                    "rope_theta": 1000,
                    "use_qk_norm": True,
                    "vocab_size": 1000,
                },
                "num_hidden_layers": 4,
                "vocab_size": 1000,
            },
            {
                "model_type": "longcat_flash_ngram",
                "attention_method": "MLA",
                "zero_expert_type": "identity",
                "hidden_size": 128,
                "ffn_hidden_size": 128,
                "moe_topk": 2,
                "expert_ffn_hidden_size": 128,
                "n_routed_experts": 2,
                "zero_expert_num": 2,
                "num_layers": 4,
                "num_hidden_layers": 4,
                "vocab_size": 1000,
                "max_position_embeddings": 1000,
                "num_attention_heads": 4,
                "kv_lora_rank": 16,
                "q_lora_rank": 16,
                "qk_rope_head_dim": 8,
                "qk_nope_head_dim": 8,
                "v_head_dim": 8,
                "routed_scaling_factor": 1.0,
                "rms_norm_eps": 1e-5,
                "rope_theta": 1000,
                "mla_scale_q_lora": True,
                "mla_scale_kv_lora": True,
                "attention_bias": False,
            },
            {
                "model_type": "longcat_flash",
                "attention_method": "MLA",
                "zero_expert_type": "identity",
                "hidden_size": 128,
                "ffn_hidden_size": 128,
                "moe_topk": 2,
                "expert_ffn_hidden_size": 128,
                "n_routed_experts": 2,
                "zero_expert_num": 2,
                "num_layers": 4,
                "num_hidden_layers": 4,
                "vocab_size": 1000,
                "max_position_embeddings": 1000,
                "num_attention_heads": 4,
                "kv_lora_rank": 16,
                "q_lora_rank": 16,
                "qk_rope_head_dim": 8,
                "qk_nope_head_dim": 8,
                "v_head_dim": 8,
                "routed_scaling_factor": 1.0,
                "rms_norm_eps": 1e-5,
                "rope_theta": 1000,
                "mla_scale_q_lora": True,
                "mla_scale_kv_lora": True,
                "attention_bias": False,
            },
            {
                "model_type": "mimo",
                "hidden_size": 128,
                "num_hidden_layers": 4,
                "intermediate_size": 128,
                "num_attention_heads": 4,
                "rms_norm_eps": 1e-5,
                "vocab_size": 1000,
                "num_key_value_heads": 2,
            },
            {
                "model_type": "nemotron-nas",
                "hidden_size": 256,
                "num_hidden_layers": 4,
                "num_attention_heads": 8,
                "rms_norm_eps": 1e-5,
                "vocab_size": 128256,
                "block_configs": [
                    {
                        "attention": {
                            "n_heads_in_group": 8,
                            "no_op": False,
                            "replace_with_linear": False,
                        },
                        "ffn": {
                            "ffn_mult": 1.3125,
                            "no_op": False,
                            "replace_with_linear": False,
                        },
                    },
                ]
                * 4,
            },
            {
                "model_type": "nemotron_h",
                "vocab_size": 1000,
                "hidden_size": 128,
                "intermediate_size": 128,
                "num_hidden_layers": 4,
                "max_position_embeddings": 1000,
                "num_attention_heads": 8,
                "num_key_value_heads": 4,
                "attention_bias": False,
                "mamba_num_heads": 8,
                "mamba_head_dim": 64,
                "mamba_proj_bias": False,
                "ssm_state_size": 128,
                "conv_kernel": 3,
                "n_groups": 4,
                "time_step_limit": [1.0, 2.0],
                "mlp_bias": False,
                "layer_norm_epsilon": 1e-4,
                "rms_norm_eps": 1e-5,
                "use_bias": True,
                "use_conv_bias": True,
                "residual_in_fp32": True,
                "hybrid_override_pattern": ["*", "M", "*", "M"],
            },
            {
                "model_type": "olmoe",
                "hidden_size": 128,
                "num_hidden_layers": 4,
                "intermediate_size": 128,
                "num_attention_heads": 2,
                "rms_norm_eps": 1e-5,
                "vocab_size": 1000,
                "num_experts": 4,
                "num_experts_per_tok": 2,
            },
            {
                "model_type": "pixtral",
                "text_config": {
                    "model_type": "llama",
                    "hidden_size": 1024,
                    "num_hidden_layers": 4,
                    "intermediate_size": 2048,
                    "num_attention_heads": 4,
                    "rms_norm_eps": 1e-5,
                    "vocab_size": 1000,
                },
                "num_hidden_layers": 4,
                "vocab_size": 1000,
            },
            {
                "model_type": "qwen3_vl_moe",
                "text_config": {
                    "model_type": "qwen3_moe",
                    "hidden_size": 128,
                    "num_hidden_layers": 4,
                    "intermediate_size": 256,
                    "num_attention_heads": 4,
                    "num_key_value_heads": 2,
                    "rms_norm_eps": 1e-5,
                    "head_dim": 32,
                    "vocab_size": 1000,
                    "decoder_sparse_step": 1,
                    "mlp_only_layers": [],
                    "num_experts_per_tok": 2,
                    "num_experts": 4,
                    "moe_intermediate_size": 128,
                    "rope_theta": 1000,
                    "max_position_embeddings": 1000,
                    "tie_word_embeddings": False,
                    "norm_topk_prob": True,
                },
                "num_hidden_layers": 4,
                "vocab_size": 1000,
            },
            {
                "model_type": "qwen3_vl",
                "text_config": {
                    "model_type": "qwen3",
                    "hidden_size": 128,
                    "num_hidden_layers": 4,
                    "intermediate_size": 256,
                    "num_attention_heads": 4,
                    "num_key_value_heads": 2,
                    "rms_norm_eps": 1e-5,
                    "vocab_size": 1000,
                    "head_dim": 32,
                    "max_position_embeddings": 1000,
                    "tie_word_embeddings": False,
                    "rope_theta": 1000,
                },
                "num_hidden_layers": 4,
                "vocab_size": 1000,
            },
            {
                "model_type": "seed_oss",
                "hidden_size": 128,
                "num_hidden_layers": 4,
                "intermediate_size": 128,
                "num_attention_heads": 2,
                "rms_norm_eps": 1e-5,
                "vocab_size": 1000,
                "num_key_value_heads": 2,
                "head_dim": 64,
            },
            {
                "model_type": "Klear",
                "hidden_size": 128,
                "num_hidden_layers": 4,
                "intermediate_size": 128,
                "num_attention_heads": 4,
                "attention_bias": False,
                "mlp_only_layers": [0],
                "num_experts": 4,
                "num_experts_per_tok": 2,
                "decoder_sparse_step": 2,
                "n_shared_experts": 1,
                "moe_intermediate_size": 128,
                "rms_norm_eps": 1e-5,
                "vocab_size": 1000,
                "num_key_value_heads": 4,
                "rope_theta": 1000.0,
                "max_position_embeddings": 1000,
                "norm_topk_prob": True,
            },
            {
                "model_type": "lille-130m",
                "block_size": 128,
                "num_hidden_layers": 4,
                "n_layer": 4,
                "n_head": 4,
                "n_kv_heads": 4,
                "n_embd": 128,
                "vocab_size": 1000,
                "rope_theta": 1000,
                "layer_norm_eps": 1e-5,
            },
            {
                "model_type": "granitemoehybrid",
                "vocab_size": 1000,
                "hidden_size": 128,
                "intermediate_size": 128,
                "num_hidden_layers": 4,
                "max_position_embeddings": 1000,
                "num_attention_heads": 8,
                "num_key_value_heads": 4,
                "attention_bias": False,
                "embedding_multiplier": 1.0,
                "attention_multiplier": 1.0,
                "logits_scaling": 1.0,
                "residual_multiplier": 1.0,
                "num_local_experts": 8,
                "num_experts_per_tok": 2,
                "shared_intermediate_size": 128,
                "mamba_n_heads": 8,
                "mamba_d_head": 16,
                "mamba_proj_bias": False,
                "mamba_d_state": 128,
                "mamba_d_conv": 4,
                "mamba_n_groups": 1,
                "mamba_conv_bias": False,
                "layer_types": ["mamba", "attention", "mamba", "attention"],
                "rms_norm_eps": 1e-5,
                "rope_theta": 1000.0,
            },
            {
                "model_type": "glm",
                "hidden_size": 128,
                "num_hidden_layers": 4,
                "intermediate_size": 128,
                "num_attention_heads": 4,
                "rms_norm_eps": 1e-5,
                "vocab_size": 1000,
                "head_dim": 32,
                "num_key_value_heads": 2,
            },
            {
                "model_type": "llama4_text",
                "hidden_size": 128,
                "num_hidden_layers": 4,
                "intermediate_size": 128,
                "num_attention_heads": 4,
                "rms_norm_eps": 1e-5,
                "vocab_size": 1000,
                "head_dim": 32,
                "num_key_value_heads": 2,
                "intermediate_size_mlp": 128,
                "rope_theta": 1000.0,
                "head_dim": 8,
                "tie_word_embeddings": False,
                "no_rope_layers": [0, 0, 1, 1],
                "use_qk_norm": True,
            },
            {
                "model_type": "mamba2",
                "num_heads": 8,
                "head_dim": 16,
                "vocab_size": 1000,
                "hidden_size": 128,
                "intermediate_size": 128,
                "state_size": 32,
                "num_hidden_layers": 4,
                "layer_norm_epsilon": 1e-4,
                "conv_kernel": 3,
                "n_groups": 4,
                "use_bias": False,
                "use_conv_bias": False,
                "tie_word_embeddings": True,
                "time_step_limit": (0.01, 10),
                "time_step_rank": "auto",
            },
            {
                "model_type": "olmo3",
                "num_heads": 8,
                "head_dim": 16,
                "vocab_size": 1000,
                "hidden_size": 128,
                "intermediate_size": 128,
                "num_attention_heads": 8,
                "rope_theta": 1000,
                "num_hidden_layers": 8,
                "rms_norm_eps": 1e-4,
                "sliding_window": 128,
                "tie_word_embeddings": True,
                "max_position_embeddings": 1000,
            },
            {
                "model_type": "jamba",
                "hidden_size": 128,
                "intermediate_size": 128,
                "num_hidden_layers": 8,
                "num_attention_heads": 4,
                "num_key_value_heads": 2,
                "attn_layer_offset": 1,
                "attn_layer_period": 2,
                "expert_layer_offset": 1,
                "expert_layer_period": 2,
                "mamba_d_conv": 4,
                "mamba_d_state": 128,
                "mamba_expand": 128,
                "num_experts": 4,
                "num_experts_per_tok": 2,
                "rms_norm_eps": 1e-5,
                "max_position_embeddings": 1000,
                "vocab_size": 1000,
            },
            {
                "model_type": "nanochat",
                "hidden_size": 1280,
                "num_hidden_layers": 20,
                "vocab_size": 32,
                "intermediate_size": 128,
            },
            {
                "model_type": "minimax",
                "hidden_size": 128,
                "intermediate_size": 128,
                "num_attention_heads": 8,
                "num_key_value_heads": 8,
                "max_position_embeddings": 1000,
                "num_experts_per_tok": 2,
                "num_local_experts": 8,
                "shared_intermediate_size": 128,
                "num_hidden_layers": 4,
                "rms_norm_eps": 1e-4,
                "rope_theta": 1000,
                "rotary_dim": 16,
                "vocab_size": 1000,
            },
            {
                "model_type": "bailing_moe_linear",
                "hidden_size": 1024,
                "num_hidden_layers": 4,
                "intermediate_size": 2048,
                "moe_intermediate_size": 1024,
                "num_experts_per_tok": 2,
                "num_experts": 4,
                "norm_topk_prob": True,
                "num_shared_experts": 2,
                "num_attention_heads": 4,
                "num_key_value_heads": 4,
                "rms_norm_eps": 1e-5,
                "vocab_size": 10_000,
                "rope_theta": 1000,
                "first_k_dense_replace": 0,
                "layer_group_size": 2,
                "group_norm_size": 1,
                "max_position_embeddings": 1000,
            },
            {
                "model_type": "qwen3_next",
                "hidden_size": 128,
                "num_hidden_layers": 4,
                "intermediate_size": 128,
                "num_attention_heads": 8,
                "num_key_value_heads": 4,
                "vocab_size": 1000,
                "linear_num_value_heads": 4,
                "linear_num_key_heads": 4,
                "linear_key_head_dim": 32,
                "linear_value_head_dim": 32,
                "linear_conv_kernel_dim": 3,
                "num_experts": 4,
                "num_experts_per_tok": 2,
                "decoder_sparse_step": 1,
                "shared_expert_intermediate_size": 128,
                "mlp_only_layers": [0],
                "moe_intermediate_size": 128,
                "rms_norm_eps": 1e-5,
                "head_dim": 64,
                "rope_theta": 1000.0,
                "partial_rotary_factor": 0.5,
                "max_position_embeddings": 1000,
            },
            {
                "model_type": "qwen3_5",
                "hidden_size": 128,
                "num_hidden_layers": 4,
                "intermediate_size": 128,
                "num_attention_heads": 8,
                "num_key_value_heads": 4,
                "vocab_size": 1000,
                "linear_num_value_heads": 4,
                "linear_num_key_heads": 4,
                "linear_key_head_dim": 32,
                "linear_value_head_dim": 32,
                "linear_conv_kernel_dim": 3,
                "rms_norm_eps": 1e-5,
                "head_dim": 64,
                "rope_theta": 1000.0,
                "partial_rotary_factor": 0.5,
                "max_position_embeddings": 1000,
            },
            {
                "model_type": "qwen3_5_moe",
                "hidden_size": 128,
                "num_hidden_layers": 4,
                "num_attention_heads": 8,
                "num_key_value_heads": 4,
                "vocab_size": 1000,
                "linear_num_value_heads": 4,
                "linear_num_key_heads": 4,
                "linear_key_head_dim": 32,
                "linear_value_head_dim": 32,
                "linear_conv_kernel_dim": 3,
                "num_experts": 4,
                "num_experts_per_tok": 2,
                "shared_expert_intermediate_size": 128,
                "moe_intermediate_size": 128,
                "rms_norm_eps": 1e-5,
                "head_dim": 64,
                "rope_theta": 1000.0,
                "partial_rotary_factor": 0.5,
                "max_position_embeddings": 1000,
            },
            {
                "model_type": "kimi_linear",
                "vocab_size": 1000,
                "hidden_size": 128,
                "num_hidden_layers": 4,
                "num_attention_heads": 8,
                "num_key_value_heads": 4,
                "intermediate_size": 128,
                "head_dim": 32,
                "rope_theta": 100.0,
                "rms_norm_eps": 1e-6,
                "linear_attn_config": {
                    "num_heads": 8,
                    "head_dim": 32,
                    "kda_layers": [1],
                },
                "model_max_length": 1000,
                "num_experts": 2,
                "moe_intermediate_size": 128,
                "kv_lora_rank": 8,
                "qk_nope_head_dim": 16,
                "qk_rope_head_dim": 16,
                "v_head_dim": 16,
            },
            {
                "model_type": "afmoe",
                "vocab_size": 1000,
                "hidden_size": 128,
                "num_hidden_layers": 4,
                "num_attention_heads": 8,
                "num_key_value_heads": 4,
                "intermediate_size": 128,
                "head_dim": 32,
                "rope_theta": 100.0,
                "layer_types": [
                    "full_attention",
                    "sliding_attention",
                    "sliding_attention",
                    "full_attention",
                ],
                "num_experts": 4,
                "num_experts_per_tok": 2,
                "moe_intermediate_size": 128,
            },
            {
                "model_type": "deepseek_v32",
                "vocab_size": 1024,
                "hidden_size": 128,
                "intermediate_size": 256,
                "moe_intermediate_size": 256,
                "num_hidden_layers": 4,
                "num_attention_heads": 4,
                "num_key_value_heads": 2,
                "n_routed_experts": 4,
                "n_group": 2,
                "topk_group": 1,
                "num_experts_per_tok": 2,
                "n_shared_experts": 1,
                "kv_lora_rank": 4,
                "q_lora_rank": 4,
                "qk_rope_head_dim": 32,
                "v_head_dim": 16,
                "qk_nope_head_dim": 32,
                "rope_scaling": {
                    "beta_fast": 32,
                    "beta_slow": 1,
                    "factor": 40,
                    "mscale": 1.0,
                    "mscale_all_dim": 1.0,
                    "original_max_position_embeddings": 4096,
                    "type": "yarn",
                },
            },
            {
                "model_type": "mimo_v2_flash",
                "num_experts_per_tok": 2,
                "hybrid_layer_pattern": [0, 1, 0, 1],
                "moe_layer_freq": [0, 1, 0, 1],
                "add_swa_attention_sink_bias": True,
                "add_full_attention_sink_bias": False,
                "sliding_window_size": 32,
                "vocab_size": 1000,
                "hidden_size": 512,
                "intermediate_size": 512,
                "moe_intermediate_size": 128,
                "num_hidden_layers": 4,
                "num_attention_heads": 4,
                "num_key_value_heads": 2,
                "n_shared_experts": 1,
                "n_routed_experts": 8,
                "routed_scaling_factor": None,
                "topk_method": "noaux_tc",
                "scoring_func": "sigmoid",
                "norm_topk_prob": True,
                "n_group": 2,
                "topk_group": 1,
                "max_position_embeddings": 1000,
                "layernorm_epsilon": 1e-5,
                "rope_theta": 1000.0,
                "swa_rope_theta": 1000.0,
                "swa_num_attention_heads": 4,
                "swa_num_key_value_heads": 2,
                "head_dim": 128,
                "v_head_dim": 64,
                "swa_head_dim": 128,
                "swa_v_head_dim": 64,
                "partial_rotary_factor": 0.5,
            },
            {
                "model_type": "rwkv7",
                "vocab_size": 1000,
                "hidden_size": 128,
                "intermediate_size": 128,
                "norm_eps": 1e-5,
                "head_dim": 32,
                "num_hidden_layers": 4,
                "a_low_rank_dim": 16,
                "v_low_rank_dim": 16,
                "gate_low_rank_dim": 16,
                "decay_low_rank_dim": 16,
            },
            {
                "model_type": "exaone_moe",
                "vocab_size": 1000,
                "hidden_size": 128,
                "intermediate_size": 256,
                "moe_intermediate_size": 64,
                "num_hidden_layers": 4,
                "num_attention_heads": 4,
                "num_key_value_heads": 2,
                "head_dim": 32,
                "num_experts": 4,
                "num_experts_per_tok": 2,
                "num_shared_experts": 1,
                "n_group": 1,
                "topk_group": 1,
                "routed_scaling_factor": 2.5,
                "norm_topk_prob": True,
                "sliding_window": 32,
                "max_position_embeddings": 1000,
                "rms_norm_eps": 1e-5,
                "rope_theta": 1000.0,
                "layer_types": [
                    "sliding_attention",
                    "sliding_attention",
                    "sliding_attention",
                    "full_attention",
                ],
                "is_moe_layer": [False, True, True, True],
                "tie_word_embeddings": False,
            },
            {
                "model_type": "youtu_llm",
                "vocab_size": 1000,
                "hidden_size": 128,
                "intermediate_size": 128,
                "num_hidden_layers": 4,
                "kv_lora_rank": 128,
                "q_lora_rank": 256,
            },
            {
                "model_type": "telechat3",
                "hidden_size": 64,
                "num_hidden_layers": 4,
                "intermediate_size": 256,
                "num_attention_heads": 8,
                "num_key_value_heads": 4,
                "rms_norm_eps": 1e-5,
                "vocab_size": 128,
                "rope_theta": 10000.0,
                "max_position_embeddings": 1000,
            },
        ]
        for config in test_configs:
            model_type = config["model_type"]
            with self.subTest(f"Testing {model_type}", model_type=model_type):
                arch = importlib.import_module(f"mlx_lm.models.{model_type}")
                args = arch.ModelArgs.from_dict(config)
                model = arch.Model(args)
                self.model_test_runner(
                    model,
                    args.model_type,
                    config["vocab_size"],
                    config["num_hidden_layers"],
                )

    def test_ssm(self):
        for batch_size in [1, 2]:
            for n_group in [1, 4]:
                num_heads = 48
                head_dim = 64
                state_dim = 128

                hidden_states = mx.random.normal(
                    shape=(batch_size, 1, num_heads, head_dim)
                )
                B = mx.random.normal(shape=(batch_size, 1, n_group, state_dim))
                C = mx.random.normal(shape=(batch_size, 1, n_group, state_dim))
                dt = mx.random.normal(shape=(batch_size, 1, num_heads))
                dt_bias = mx.random.normal(shape=(num_heads,))
                A_log = mx.random.normal(shape=(num_heads,))
                D = mx.random.normal(shape=(num_heads,))
                state = mx.random.normal(
                    shape=(batch_size, num_heads, head_dim, state_dim)
                )

                out, out_state = ssm_attn(
                    hidden_states, A_log, B, C, D, dt, dt_bias, state
                )
                out_c, out_state_c = ssm_update(
                    hidden_states, A_log, B, C, D, dt, dt_bias, state
                )
                self.assertTrue(mx.allclose(out, out_c, atol=1e-4, rtol=1e-4))
                self.assertTrue(
                    mx.allclose(out_state, out_state_c, atol=1e-4, rtol=1e-4)
                )

    def test_ssm_masked(self):
        batch_size = 1
        n_group = 1
        num_heads = 48
        head_dim = 64
        state_dim = 128
        seq_len = 4
        pad = 2

        hidden_states = mx.random.normal(
            shape=(batch_size, seq_len + pad, num_heads, head_dim)
        )
        B = mx.random.normal(shape=(batch_size, seq_len + pad, n_group, state_dim))
        C = mx.random.normal(shape=(batch_size, seq_len + pad, n_group, state_dim))
        dt = mx.random.normal(shape=(batch_size, seq_len + pad, num_heads))
        dt_bias = mx.random.normal(shape=(num_heads,))
        A_log = mx.random.normal(shape=(num_heads,))
        D = mx.random.normal(shape=(num_heads,))
        out, out_state = ssm_attn(
            hidden_states[:, pad:],
            A_log,
            B[:, pad:],
            C[:, pad:],
            D,
            dt[:, pad:],
            dt_bias,
        )
        mask = mx.array([[False] * pad + [True] * seq_len])
        out_m, out_state_m = ssm_attn(
            hidden_states, A_log, B, C, D, dt, dt_bias, mask=mask
        )
        out_m = out_m[:, pad:]
        self.assertTrue(mx.allclose(out, out_m, atol=1e-4, rtol=1e-4))
        self.assertTrue(mx.allclose(out_state, out_state_m, atol=1e-4, rtol=1e-4))

    def test_ssm_right_pad(self):
        batch_size = 1
        n_group = 1
        num_heads = 48
        head_dim = 64
        state_dim = 128
        seq_len = 4
        pad = 2

        hidden_states = mx.random.normal(
            shape=(batch_size, seq_len + pad, num_heads, head_dim)
        )
        B = mx.random.normal(shape=(batch_size, seq_len + pad, n_group, state_dim))
        C = mx.random.normal(shape=(batch_size, seq_len + pad, n_group, state_dim))
        dt = mx.random.normal(shape=(batch_size, seq_len + pad, num_heads))
        dt_bias = mx.random.normal(shape=(num_heads,))
        A_log = mx.random.normal(shape=(num_heads,))
        D = mx.random.normal(shape=(num_heads,))
        out, out_state = ssm_attn(
            hidden_states[:, :-pad],
            A_log,
            B[:, :-pad],
            C[:, :-pad],
            D,
            dt[:, :-pad],
            dt_bias,
        )
        mask = mx.array([[True] * seq_len + [False] * pad])
        lengths = mx.array([seq_len])
        out_m, out_state_m = ssm_attn(
            hidden_states,
            A_log,
            B,
            C,
            D,
            dt,
            dt_bias,
            mask=mask,
            lengths=lengths,
        )
        out_m = out_m[:, :-pad]
        self.assertTrue(mx.allclose(out, out_m, atol=1e-4, rtol=1e-4))
        self.assertTrue(mx.allclose(out_state, out_state_m, atol=1e-4, rtol=1e-4))

    def test_gated_delta(self):
        mx.random.seed(0)
        for B in [1, 2]:
            for T in [1, 2]:
                Hk = 16
                Hv = 32
                Dk = 128
                Dv = 128

                q = mx.random.normal(shape=(B, T, Hk, Dk))
                k = mx.random.normal(shape=(B, T, Hk, Dk))
                v = mx.random.normal(shape=(B, T, Hv, Dv))
                g = mx.random.uniform(shape=(B, T, Hv))
                beta = mx.random.uniform(shape=(B, T, Hv))
                state = mx.random.normal(shape=(B, Hv, Dk, Dv))

                y_op, st_op = gated_delta_ops(q, k, v, g, beta, state)
                y_c, st_c = gated_delta_kernel(q, k, v, g, beta, state)
                self.assertTrue(mx.allclose(y_op, y_c, rtol=1e-4, atol=1e-4))
                self.assertTrue(mx.allclose(st_op, st_c, rtol=1e-4, atol=1e-4))

    def test_gated_delta_precision(self):
        mx.random.seed(42)

        N_STEPS = 512
        B = 1
        Hk = 4
        Hv = 4
        Dk = 64
        Dv = 64

        A_log = mx.zeros((Hv,))
        dt_bias = mx.ones((Hv,))

        all_q = mx.random.normal(shape=(N_STEPS, B, 1, Hk, Dk)) * 0.1
        all_k = mx.random.normal(shape=(N_STEPS, B, 1, Hk, Dk)) * 0.1
        all_v = mx.random.normal(shape=(N_STEPS, B, 1, Hv, Dv)) * 0.1
        all_a = -7.0 + mx.random.normal(shape=(N_STEPS, B, 1, Hv)) * 0.3
        all_b = mx.random.normal(shape=(N_STEPS, B, 1, Hv))
        mx.eval(all_q, all_k, all_v, all_a, all_b, A_log, dt_bias)

        state_ref = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32)
        for t in range(N_STEPS):
            y_ref, state_ref = gated_delta_update(
                all_q[t],
                all_k[t],
                all_v[t],
                all_a[t],
                all_b[t],
                A_log,
                dt_bias,
                state_ref,
                use_kernel=False,
            )
            mx.eval(y_ref, state_ref)

        for use_kernel in (False, True):
            state_lo = mx.zeros((B, Hv, Dv, Dk), dtype=mx.bfloat16)
            for t in range(N_STEPS):
                y_lo, state_lo = gated_delta_update(
                    all_q[t].astype(mx.bfloat16),
                    all_k[t].astype(mx.bfloat16),
                    all_v[t].astype(mx.bfloat16),
                    all_a[t].astype(mx.bfloat16),
                    all_b[t].astype(mx.bfloat16),
                    A_log,
                    dt_bias,
                    state_lo,
                    use_kernel=use_kernel,
                )
                mx.eval(y_lo, state_lo)

            self.assertTrue(mx.allclose(state_lo, state_ref, rtol=0.05, atol=0.01))
            self.assertTrue(mx.allclose(y_lo, y_ref, rtol=0.05, atol=0.01))

    def test_gated_delta_masked(self):
        B = 1
        T = 3
        Hk = 16
        Hv = 32
        Dk = 128
        Dv = 128

        mx.random.seed(0)
        q = mx.random.normal(shape=(B, T, Hk, Dk))
        k = mx.random.normal(shape=(B, T, Hk, Dk))
        v = mx.random.normal(shape=(B, T, Hv, Dv))
        g = mx.random.normal(shape=(B, T, Hv))
        beta = mx.random.normal(shape=(B, T, Hv))
        state = mx.random.normal(shape=(B, Hv, Dk, Dv))

        for s, e, mask in [
            (1, 3, mx.array([[False, True, True]])),
            (0, 2, mx.array([[True, True, False]])),
        ]:
            y_gt, st_gt = gated_delta_ops(
                q[:, s:e],
                k[:, s:e],
                v[:, s:e],
                g[:, s:e],
                beta[:, s:e],
                state,
            )
            for fn in [gated_delta_ops, gated_delta_kernel]:
                y, st = fn(q, k, v, g, beta, state, mask)
                y = y[:, s:e]
                self.assertTrue(mx.allclose(y, y_gt, rtol=1e-4, atol=1e-4))
                self.assertTrue(mx.allclose(st, st_gt, rtol=1e-4, atol=1e-3))


if __name__ == "__main__":
    unittest.main()