"""Real, tiny hybrid-model regressions; no network or large weights required.""" import copy import platform import pytest if platform.system() != "Darwin" or platform.machine() != "arm64": pytest.skip("MLX requires Apple Silicon", allow_module_level=True) mx = pytest.importorskip("mlx.core") pytest.importorskip("mlx_lm") from mlx_lm.models.qwen3_5 import Model, ModelArgs from semif_phase1 import mlx_backend as backend class Tokenizer: pad_token_id = 0 eos_token_id = 1 def apply_chat_template(self, turns, **kwargs): return "\n".join(turn["content"] for turn in turns) + "\nAssistant:" def encode(self, text, add_special_tokens=False): return list(text.encode()) def decode(self, ids): return bytes(ids).decode() @pytest.fixture(scope="module") def model(): mx.random.seed(17) result = Model(ModelArgs(model_type="qwen3_5", text_config={ "hidden_size": 64, "intermediate_size": 128, "num_hidden_layers": 4, "num_attention_heads": 4, "num_key_value_heads": 2, "head_dim": 16, "vocab_size": 256, "linear_num_key_heads": 2, "linear_num_value_heads": 4, # The native Metal delta kernel requires production-sized head widths. "linear_key_head_dim": 128, "linear_value_head_dim": 128, "full_attention_interval": 2, })) result.eval() mx.eval(result.parameters()) return result @pytest.fixture def rows(): return [dict(id=str(i), state={"evidence": "A deployment succeeded."}, question=question, options=[{"id": "yes", "description": "Yes"}, {"id": "no", "description": "No"}]) for i, question in enumerate(["Success?", "Did the deployment succeed in the supplied evidence?", "Failure?"])] def assert_same(left, right): assert left["id"] == right["id"] for key in ("option_ids", "answer_token_ids", "input_ids_sha256", "prompt_sha256"): assert left[key] == right[key] # Full and split prefill use different Metal reduction shapes, even in FP32. assert left["probabilities"] == pytest.approx(right["probabilities"], abs=1e-4) def test_hybrid_cache_branches_padding_and_order(model, rows): tokenizer = Tokenizer() fresh = [backend.score(model, tokenizer, row, {}) for row in rows] serial = backend.SerialPrefixScorer(model, tokenizer, {}) for index in (0, 2, 1, 0): actual = serial.score(rows[index]) assert_same(fresh[index], actual) for order in (rows, rows[::-1], rows[:1]): results, timing = backend.score_shared(model, tokenizer, order, {}) assert timing["batch_size"] == len(order) for row in results: assert_same(fresh[int(row["id"])], row) def test_serial_invalidates_cache_for_mutated_structured_state(model, rows): row = copy.deepcopy(rows[0]) serial = backend.SerialPrefixScorer(model, Tokenizer(), {}) assert not serial.score(row)["cache_hit"] assert serial.score(row)["cache_hit"] row["state"]["evidence"] = "The deployment failed." actual = serial.score(row) assert not actual["cache_hit"] assert_same(backend.score(model, Tokenizer(), row, {}), actual) def test_shared_rejects_mixed_states_and_duplicate_ids(model, rows): with pytest.raises(ValueError, match="unique"): backend.score_shared(model, Tokenizer(), [rows[0], rows[0]], {}) rows[1]["state"] = "different" with pytest.raises(ValueError, match="exact state"): backend.score_shared(model, Tokenizer(), rows, {}) def test_token_limit_prevents_inference(model, rows): with pytest.raises(ValueError, match="no truncation"): backend.score(model, Tokenizer(), rows[0], {}, max_tokens=2) @pytest.mark.parametrize("config", [ {"model_type": "qwen3_5", "model_file": "custom.py"}, {"model_type": "unrecognized"}, ]) def test_loader_rejects_custom_or_unsupported_models(tmp_path, config): import json (tmp_path / "config.json").write_text(json.dumps(config)) with pytest.raises(ValueError, match="native Qwen3.5"): backend.load_model(str(tmp_path), "local-fixture") def test_remote_model_requires_immutable_revision(): with pytest.raises(ValueError, match="40-character"): backend.load_model("Qwen/Qwen3.5-4B", "main") def test_recurrent_qk_normalization_matches_reference_l2_epsilon(): # Small q/k expose the release-0.31.3 sum-versus-mean epsilon mismatch. from mlx_lm.models.gated_delta import normalize_qk width = 128 values = mx.full((1, 2, width), 1e-4) q, k = normalize_qk(values, values, inv_scale=width**-0.5, eps=1e-6) expected = values * mx.rsqrt(mx.sum(values * values, axis=-1, keepdims=True) + 1e-6) assert mx.max(mx.abs(k - expected)).item() < 1e-6 assert mx.max(mx.abs(q - expected * width**-0.5)).item() < 1e-6 def test_serial_keys_by_exact_tokens_not_python_value_equality(model, rows): row = copy.deepcopy(rows[0]) row["state"] = {"value": True} serial = backend.SerialPrefixScorer(model, Tokenizer(), {}) serial.score(row) row["state"] = {"value": 1} # Python compares True == 1, JSON prompts differ. actual = serial.score(row) assert not actual["cache_hit"] assert_same(backend.score(model, Tokenizer(), row, {}), actual) @pytest.mark.parametrize('limit_mib', [None, 0, 64]) def test_loader_applies_cache_limit_and_records_bytes(tmp_path, monkeypatch, model, limit_mib): import json import mlx_lm (tmp_path / 'config.json').write_text(json.dumps({'model_type': 'qwen3_5'})) monkeypatch.setattr(mlx_lm, 'load', lambda *args, **kwargs: (model, Tokenizer())) previous = mx.set_cache_limit(32 * 1024 * 1024) try: kwargs = {} if limit_mib is None else {'cache_limit_mib': limit_mib} _, _, metadata = backend.load_model(str(tmp_path), 'local-fixture', **kwargs) expected = (256 if limit_mib is None else limit_mib) * 1024 * 1024 assert metadata['allocator_cache_limit_bytes'] == expected assert mx.set_cache_limit(previous) == expected finally: mx.set_cache_limit(previous) @pytest.mark.parametrize('limit', [-1, 1.5, True]) def test_loader_rejects_invalid_cache_limit_before_loading(limit): with pytest.raises(ValueError, match='nonnegative integer'): backend.load_model('Qwen/Qwen3.5-4B', '0' * 40, cache_limit_mib=limit)