Add test manifest for PR #1410 (glm_moe_dsa)

harness/manifest-schema.yaml

@@ -37,6 +37,9 @@ variants:
     # Expected output dtype. Read from config by the skill.
     # The harness verifies both forward-pass output and norm weight dtypes.
     expected_dtype: bfloat16
+    # Optional: a different HF repo to download and run the MLX tests against.
+    # Useful when the canonical model is too big and you want to test a quant.
+    # mlx_test_repo: pcuenq/Olmo-Hybrid-7B-mxfp4
 
   - repo_id: allenai/Olmo-Hybrid-Instruct-DPO-7B
     type: instruct

harness/run_harness.py

@@ -248,12 +248,9 @@ def bucket_has_reference(repo_id):
     )
     if r.returncode != 0:
         return False
-    # Each line is "<size> [<date>] <name>". Check the basename column for
-    # an exact match — substring match would false-positive on siblings
-    # like `numerical.safetensors.bak`.
     for line in r.stdout.splitlines():
         parts = line.split()
-        if parts and parts[-1] == "numerical.safetensors":
+        if parts and parts[-1].rsplit("/", 1)[-1] == "numerical.safetensors":
             return True
     return False
 
@@ -356,9 +353,14 @@ def run_test(name, args, output_dir, timeout=300):
 
 
 def run_battery(variant, python, output_dir, work_dir, quantize,
-                trust_remote_code=False, reference_path=None):
+                trust_remote_code=False, reference_path=None,
+                model_path=None):
+    """`model_path` overrides the model location passed to all MLX CLI / test
+    invocations. Defaults to `variant["repo_id"]`. Set by callers when the
+    manifest has `mlx_test_repo` and we should download/test against that
+    instead of the canonical source."""
     vtype = variant["type"]
-    model = variant["repo_id"]
+    model = model_path or variant["repo_id"]
 
     gen_prompt = BASE_GENERATE_PROMPT if vtype == "base" else INSTRUCT_GENERATE_PROMPT
     long_prompt = BASE_LONG_PROMPT if vtype == "base" else INSTRUCT_LONG_PROMPT
@@ -811,8 +813,13 @@ def main():
             output_dir = results_dir / slug
             output_dir.mkdir(parents=True, exist_ok=True)
 
-            log(f"=== {variant['repo_id']} ({variant['type']}) ===")
-            download_model(variant["repo_id"], python)
+            mlx_test_repo = variant.get("mlx_test_repo") or variant["repo_id"]
+            if mlx_test_repo != variant["repo_id"]:
+                log(f"=== {variant['repo_id']} ({variant['type']}) "
+                    f"— testing against {mlx_test_repo} ===")
+            else:
+                log(f"=== {variant['repo_id']} ({variant['type']}) ===")
+            download_model(mlx_test_repo, python)
 
             # Bucket existence pre-checked above for cuda_reference variants.
             reference_path = None
@@ -823,6 +830,7 @@ def main():
                 variant, python, output_dir, work_dir, quantize,
                 trust_remote_code=trust_remote_code,
                 reference_path=reference_path,
+                model_path=mlx_test_repo,
             )
 
         write_summary(manifest, results_dir, sha)

manifests/pr-1410.yaml

@@ -0,0 +1,41 @@
+version: 1
+pr:
+  number: 1410
+  repo: pcuenca/mlx-lm
+  branch: glm-moe-dsa-indexer-sharing
+model_type: glm_moe_dsa
+variants:
+  - repo_id: zai-org/GLM-5.2-FP8
+    type: instruct
+    memory_gb: 1500
+    expected_dtype: bfloat16
+    mlx_test_repo: mlx-community/GLM-5.2-mxfp4
+quantize: false
+trust_remote_code: false
+cuda_reference: true
+extra_pip:
+  - transformers[kernels]==5.12.0
+notes: >
+  GLM-5.2 (~750B MoE, MLA + DeepSeek Sparse Attention). This PR adds the GLM-5.2
+  novelty: a per-layer DSA indexer schedule with cross-layer top-k sharing
+  (config.indexer_types — "full" layers run their own indexer; "shared" layers
+  reuse the previous full layer's top-k). Things the fixed battery may not cover:
+
+  - Converted from the FP8 checkpoint (e4m3, 128x128 block, weight_scale_inv).
+    mlx-lm dequantizes FP8 -> bf16 on load (deepseek_v32.sanitize), then quantizes
+    to mxfp4. The MLX model is ~368 GB and fits a single 512 GB machine despite the
+    ~1500 GB unquantized (bf16) figure above; memory_gb is the unquantized estimate.
+
+  - The DSA indexer only sparsifies once the sequence exceeds index_topk=2048, so
+    the numerical/per-layer comparison MUST use a >2048-token sequence to exercise
+    the sharing path. Short prompts test only the dense path and will look fine even
+    if the sparse path is wrong.
+
+  - RoPE caveat for the CUDA reference (important): GLM-5.2's indexer uses
+    INTERLEAVED RoPE (config indexer_rope_interleave: true; mlx traditional=True,
+    which this PR keeps). Current upstream transformers' glm_moe_dsa indexer appears
+    to use half-split RoPE (apply_rotary_pos_emb) — a likely regression. If the CUDA
+    reference runs stock transformers, a >2048-token comparison may show a SPURIOUS
+    mismatch where MLX is correct and the reference is wrong. Verify the reference's
+    indexer RoPE before treating a long-context divergence as an MLX bug. MLX was
+    validated here via a 16384-token generation that stayed fully coherent.

results/pr-1410/2026-06-22T195215-3cb18b5/manifest.yaml

@@ -0,0 +1,41 @@
+version: 1
+pr:
+  number: 1410
+  repo: pcuenca/mlx-lm
+  branch: glm-moe-dsa-indexer-sharing
+model_type: glm_moe_dsa
+variants:
+  - repo_id: zai-org/GLM-5.2-FP8
+    type: instruct
+    memory_gb: 1500
+    expected_dtype: bfloat16
+    mlx_test_repo: mlx-community/GLM-5.2-mxfp4
+quantize: false
+trust_remote_code: false
+cuda_reference: true
+extra_pip:
+  - transformers[kernels]==5.12.0
+notes: >
+  GLM-5.2 (~750B MoE, MLA + DeepSeek Sparse Attention). This PR adds the GLM-5.2
+  novelty: a per-layer DSA indexer schedule with cross-layer top-k sharing
+  (config.indexer_types — "full" layers run their own indexer; "shared" layers
+  reuse the previous full layer's top-k). Things the fixed battery may not cover:
+
+  - Converted from the FP8 checkpoint (e4m3, 128x128 block, weight_scale_inv).
+    mlx-lm dequantizes FP8 -> bf16 on load (deepseek_v32.sanitize), then quantizes
+    to mxfp4. The MLX model is ~368 GB and fits a single 512 GB machine despite the
+    ~1500 GB unquantized (bf16) figure above; memory_gb is the unquantized estimate.
+
+  - The DSA indexer only sparsifies once the sequence exceeds index_topk=2048, so
+    the numerical/per-layer comparison MUST use a >2048-token sequence to exercise
+    the sharing path. Short prompts test only the dense path and will look fine even
+    if the sparse path is wrong.
+
+  - RoPE caveat for the CUDA reference (important): GLM-5.2's indexer uses
+    INTERLEAVED RoPE (config indexer_rope_interleave: true; mlx traditional=True,
+    which this PR keeps). Current upstream transformers' glm_moe_dsa indexer appears
+    to use half-split RoPE (apply_rotary_pos_emb) — a likely regression. If the CUDA
+    reference runs stock transformers, a >2048-token comparison may show a SPURIOUS
+    mismatch where MLX is correct and the reference is wrong. Verify the reference's
+    indexer RoPE before treating a long-context divergence as an MLX bug. MLX was
+    validated here via a 16384-token generation that stayed fully coherent.

results/pr-1410/2026-06-22T195215-3cb18b5/reproduce.md

@@ -0,0 +1,24 @@
+# Reproducing these results
+
+**PR:** https://github.com/pcuenca/mlx-lm/pull/1410 (branch `glm-moe-dsa-indexer-sharing`)
+**Commit:** `3cb18b51892a7800678923116f5b4920a7666208`
+**Model type:** `glm_moe_dsa`
+
+The exact scripts that produced these results are archived in `scripts/`.
+The manifest used is `manifest.yaml`.
+
+## Quick reproduction
+
+```bash
+python scripts/run_harness.py manifest.yaml
+```
+
+The harness resolves the branch to its current head sha. To reproduce
+against this exact commit, check out `3cb18b51892a7800678923116f5b4920a7666208` of `pcuenca/mlx-lm` manually
+and rerun the per-test commands below.
+
+## Per-test commands
+
+Each `<variant>/<test>.json` file includes the exact command that was run,
+its exit code, stdout, and stderr. To rerun a single test, copy the
+command from there.

results/pr-1410/2026-06-22T195215-3cb18b5/scripts/check_dtype.py

@@ -0,0 +1,45 @@
+#!/usr/bin/env python3
+"""Verify that a model's forward-pass output dtype matches the expected dtype."""
+import argparse
+import json
+
+import mlx.core as mx
+from mlx_lm import load
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Verify a model's forward-pass output dtype.")
+    parser.add_argument("model_path")
+    parser.add_argument("expected_dtype")
+    parser.add_argument("--trust-remote-code", action="store_true",
+                        help="Pass trust_remote_code=True to mlx_lm.load's "
+                             "tokenizer config. Required for custom-code "
+                             "models — mlx_lm.utils.load (unlike sharded_load) "
+                             "doesn't default this on.")
+    args = parser.parse_args()
+
+    tokenizer_config = {"trust_remote_code": True} if args.trust_remote_code else None
+
+    expected = getattr(mx, args.expected_dtype)
+
+    model, _ = load(args.model_path, tokenizer_config=tokenizer_config)
+
+    out = model(mx.array([[1]]))
+    mx.eval(out)
+
+    result = {
+        "output_dtype": str(out.dtype),
+        "expected_dtype": str(expected),
+        "match": out.dtype == expected,
+    }
+    print(json.dumps(result, indent=2))
+
+    if not result["match"]:
+        import sys
+        print(f"FAIL: output dtype {out.dtype} != {args.expected_dtype}", file=sys.stderr)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()