1"""Stacked-layout registry/store check (runs only where a stacked checkpoint such as Qwen3-30B-A3B-4bit is cached)."""2import glob, os, pytest, mlx.core as mx3from titan.registry import load_spec4from titan.store import ExpertStore56def _dir():7 hits = glob.glob(os.path.expanduser("~/.cache/huggingface/hub/models--mlx-community--Qwen3-30B-A3B-Instruct-2507-4bit/snapshots/*"))8 return hits[0] if hits else None910@pytest.mark.skipif(_dir() is None, reason="Qwen3-30B-A3B-Instruct-2507-4bit not cached")11def test_stacked_slices_match_mx_load():12 d = _dir(); s = load_spec(d)13 assert s.layout.startswith("stacked") and s.num_layers == 48 and s.num_experts == 128 and s.top_k == 814 st = ExpertStore(s)15 for (l, e) in [(0, 0), (23, 77), (47, 127)]:16 t = st.read_expert(l, e)17 f = s.experts[(l, e)].tensors["gate.weight"].file18 ref = mx.load(f)19 for name, arr in t.items():20 m, part = name.split(".")21 r = ref[f"model.layers.{l}.mlp.switch_mlp.{m}_proj.{part}"][e]22 assert arr.shape == r.shape and arr.dtype == r.dtype, (name, arr.shape, r.shape)23 assert bool(mx.array_equal(arr, r)), (l, e, name)24 print("bytes_per_expert", s.bytes_per_expert, "expert_bytes_total_GB", s.expert_bytes_total / 1e9)25