[transformers-to-mlx skill] glm_moe_dsa: DSA cross-layer indexer sharing (GLM-5.2)

[pcuenca]

GLM-5.2 (glm_moe_dsa) adds a per-layer DSA indexer schedule with cross-layer top-k sharing that the current bare deepseek_v32 subclass doesn't implement. This PR adds the schedule + prev_topk_indices threading, gives shared layers no indexer cache (fixes a KVCache.state crash), and adds a model test. Defaults reduce to plain deepseek_v32. Validated by FP8→mxfp4 conversion + short (TP) and 16K-token (single-node) generation. Full report in the comment below.

Disclosure: produced with the transformers-to-mlx skill (AI-assisted), human-driven and reviewed.

Test model: mlx-community/GLM-5.2-mxfp4

[pcuenca]

GLM-5.2 (glm_moe_dsa) → MLX conversion report

Disclosure: this conversion and all tests/results below were produced with the
transformers-to-mlx skill (AI-assisted). A human has driven the Skill and
helped guide the decisions.

Summary

GLM-5.2 is a ~750B-parameter MoE (GlmMoeDsaForCausalLM, model_type: glm_moe_dsa)
with Multi-head Latent Attention (MLA) + DeepSeek Sparse Attention (DSA). mlx-lm
already mapped glm_moe_dsa to a bare subclass of deepseek_v3.2. GLM-5.2 adds one
architectural feature that bare subclass does not implement: a per-layer DSA
indexer schedule with cross-layer top-k sharing. This PR implements that schedule
(and a related cache fix) and adds a model test.

A model was converted from the official FP8 checkpoint to MLX mxfp4 and
validated by generation (short + long-context) on Apple Silicon. Full-model
numerical comparison against transformers has not been performed yet due to the
large scale of the model. We expect to complete this task by extracting logits and
outputs using a CUDA test harness. We substitute a module-level indexer parity test
against the real transformers indexer for the one genuinely novel, error-prone component.

Models

Variant	Repo	Format	Notes
GLM-5.2	zai-org/GLM-5.2	bf16	~750B, full precision
GLM-5.2-FP8	zai-org/GLM-5.2-FP8	block-FP8 (e4m3, 128×128, dynamic)	conversion source
GLM-5.2 MLX	mlx-community/GLM-5.2-mxfp4	mxfp4, 4.252 bpw, 368 GB, 76 shards	tested

Expected runtime dtype: bfloat16 (source: config.json dtype: "bfloat16").

Architecture

• 78 layers; MLA: kv_lora_rank=512, q_lora_rank=2048, qk_nope_head_dim=192,
  qk_rope_head_dim=64, v_head_dim=256; num_attention_heads=64.
• MoE: 256 routed experts (8 active) + 1 shared; n_group=1, topk_method=noaux_tc,
  scoring_func=sigmoid, routed_scaling_factor=2.5; first 3 layers dense
  (first_k_dense_replace=3).
• DSA indexer: index_head_dim=128, index_n_heads=32, index_topk=2048.
• RoPE: rope_theta=8e6, max_position_embeddings=1048576. Both main attention
  (rope_interleave: true) and indexer (indexer_rope_interleave: true) use
  interleaved RoPE.
• 1 MTP layer (num_nextn_predict_layers=1) — dropped on load.

Novelty: cross-layer DSA indexer sharing

config.indexer_types is a per-layer list of "full"/"shared" (78 entries for
GLM-5.2: full at layers 0,1,2 then every 4th — index_topk_freq=4,
index_skip_topk_offset=3):

• full layers run their own DSA indexer (the lightweight wq_b/wk/
  weights_proj/k_norm projections) and produce top-k token indices.
• shared layers run no indexer (no indexer weights in the checkpoint) and
  reuse the previous full layer's top-k, threaded down as prev_topk_indices.

This is the GlmMoeDsa divergence from deepseek_v3.2 in transformers, where the
attention/decoder/model thread prev_topk_indices and shared layers set
self.indexer = None.
Reference: modeling_glm_moe_dsa.py
(GlmMoeDsaAttention, GlmMoeDsaDecoderLayer, GlmMoeDsaModel).

Changes to mlx-lm

All in mlx_lm/models/glm_moe_dsa.py (previously a bare deepseek_v32 subclass):

1. ModelArgs: add indexer_types plus index_topk_pattern / index_topk_freq
   / index_skip_topk_offset, deriving the schedule in __post_init__ (mirrors the
   transformers config). Default (freq=1) → all "full", i.e. identical to plain
   deepseek_v3.2.
2. GlmMoeDsaAttention: shared layers set self.indexer = None, consume
   prev_topk_indices, and return their top-k; the indexer-cache mx.depends
   bookkeeping is gated to full layers.
3. GlmMoeDsaDecoderLayer / GlmMoeDsaModel: thread prev_topk_indices
   layer→layer.
4. make_cache: shared layers get only the main KVCache (no indexer
   KVCache). Without this, the unused indexer cache stays empty and crashes in
   KVCache.state during generation.

No changes to deepseek_v32.py. Its indexer already uses interleaved RoPE
(traditional=True), which is correct for GLM-5.2 (see RoPE note).

Testing

What we could not do

Full-model numerical / per-layer comparison vs transformers.
This is deferred to a separate test harness.

Module-level indexer parity

glm_moe_dsa.py's indexer is deepseek_v32's. We copied real weights from the
transformers GlmMoeDsaIndexer into the mlx Indexer (tiny synthetic config),
compared top-k selections, and checked the sharing schedule against the real
config.json:

[rope] transformers indexer matched by: traditional=True -> False, traditional=False -> True
[share] using real indexer_types from config (78 layers)
[share] indexer present per layer matches schedule: True
[share] shared layers carry no indexer weights: True
[share] strict load_weights OK

RoPE note (important): GLM-5.2's config.json sets indexer_rope_interleave: true
(= mlx traditional=True). mlx-lm's indexer is already traditional=True and is
correct. The parity test shows current upstream transformers matches
traditional=False (half-split / apply_rotary_pos_emb) for the indexer — i.e.
upstream transformers appears to have a regression for the GLM-5.2 indexer RoPE,
to be verified.

Test 2 — dtype (on the converted mxfp4 model)

Tensor dtype tally across the 76 converted shards:

U8 (mxfp4 scales): 992 | U32 (mxfp4 weights): 992 | BF16: 430 | F32: 75

The 75 F32 tensors are all model.layers.N.mlp.gate.e_score_correction_bias
(one per sparse MoE layer) — intentionally kept fp32 by deepseek_v32's
cast_predicate. No stray float32 contamination.

Test 1 — short generation (distributed, tensor parallel)

2× M3 Ultra (512 GB) over a Thunderbolt ring, TP:

$ mlx.launch --backend ring --hostfile hosts.json --env MLX_METAL_FAST_SYNCH=1 \
    sharded_generate.py --model GLM-5.2-mlx-mxfp4 \
    --prompt "In one sentence, what is tensor parallelism?" -m 128
...coherent reasoning trace about splitting tensors across GPUs...
Prompt: 22 tokens, 22.941 tokens-per-sec
Generation: 128 tokens, 18.242 tokens-per-sec
Peak memory: 206.562 GB         # ~half the 368 GB model per node — TP shards correctly

Test 5 — long-context generation (single node) — the key test

Single node (model fits under the 480 GB wired limit). This is the decisive test:
the DSA indexer only sparsifies past index_topk=2048, so a wrong indexer RoPE or
sharing schedule shows up as gibberish only at long context.

$ MLX_METAL_FAST_SYNCH=1 mlx_lm.generate --model GLM-5.2-mlx-mxfp4 \
    --prompt "Write a complete, polished single-file HTML and JavaScript implementation of Space Invaders, with comments explaining each section." \
    --max-tokens 16384
Prompt: 34 tokens, 5.599 tokens-per-sec
Generation: 16384 tokens, 18.081 tokens-per-sec
Peak memory: 396.755 GB

Result: fully coherent across all 16,384 tokens. The output is a complete,
well-formed HTML/CSS/JS Space Invaders game (closing </html>, 49 structural
tags, sprite bitmaps defined as 2D arrays deep in the tail at ~16K tokens), with no
gibberish and no degenerate repetition (most-repeated lines are ordinary }/CSS).
This validates the interleaved indexer RoPE and the cross-layer sharing end-to-end
on real weights, well past the 2048 sparsification boundary.

Unit test

tests/test_models.py::test_glm_moe_dsa — builds a model with an alternating
FSFSFS schedule and asserts: the derived indexer_types, that only full layers
hold an indexer, the standard forward/cache flow (model_test_runner), the sparse
path (prompt > index_topk), and cache serialization ([c.state for c in cache],
which is what crashed before the make_cache fix).

Conversion details

Converted from the FP8 checkpoint.

$ mlx_lm.convert --hf-path GLM-5.2-FP8 --mlx-path GLM-5.2-mlx-mxfp4 --q-mode mxfp4 -q
[INFO] Quantized model with 4.252 bits per weight.

(Note: converting from FP8 quantizes already-fp8-rounded weights — marginally lossier
than converting from the bf16 source, but negligible at 4-bit and the generation
quality is clearly intact.)

Errors encountered & resolved

• KVCache.state crash during generation: shared layers never populate the
  second (indexer) KVCache, so it stayed empty and crashed when
  generate_step serialized cache state. Fixed via make_cache (shared layers get
  no indexer cache). The unit test now covers this path.

Notes / learnings

• Indexer RoPE convention is per-model and must be read from config
  (indexer_rope_interleave); the short-sequence test blind spot (≤ index_topk)
  means RoPE bugs in the indexer only surface in long-context generation.
• For a model this large that fits one node when quantized, single-node generation
  is more reliable than a 2-node ring for long runs.

Environment

mlx 0.31.2; 2× Apple M3 Ultra (512 GB), macOS; conversion target mxfp4.

[pcuenca]

Setting as draft. Initial tests pass, but we want to run additional tests and review the code in more detail.

[pcuenca]

cc @Blaizzy

[kernelpool]

Tested this on a dual mac studio setup running GLM-5.2-8bit and seems to be working well! Initially, I did something similar and extended the deepseek_v32.py implementation to support the indexer sharing, which results in overall fewer changes, but I think this approach is better, since the extension is glm_moe_dsa specific.

[bibproj2]

Also tested this, but on a single mac studio. Passed all my tests for coding and translating accuracy.

[ivanfioravanti]

I will test this tomorrow! Thanks @pcuenca for conversion and converter! What an amazing job there! 🙏

[pcuenca]

Thank you @ivanfioravanti! For additional context, I'm trying to compare hidden states against the transformers implementation, but the process is giving me some trouble (infra issues). Will update when I can!

[spicyneuron]

I used this PR to quantize for a single M3 Ultra yesterday, and it's been driving my Claude Code ever since. Very stable and extremely capable model, if a bit slow. Thanks @pcuenca!

[bibproj2]

There seem to be some issues:

In https://huggingface.co/mlx-community/GLM-5.2-DQ4plus-q8/discussions/3 people mention that as the prompt gets longer, the MLX quants based on this PR become incoherent.

@pcuenca make a mxfp4 quant from GLM-5.2-FP8, which works. When I tried to make a mxfp4 quant from the BF16 version, it created it, but generate threw an error,

... model.load_weights(list(weights.items()), strict=strict)
  File "/opt/homebrew/lib/python3.11/site-packages/mlx/nn/layers/base.py", line 191, in load_weights
    raise ValueError(f"Missing {num_missing} parameters: \n{missing}.")
ValueError: Missing 1142 parameters: 
lm_head.weight, ...

[c-geek]

Fixes everything for me! (EXO with 2 Mac Studios against kernelpool/GLM-5.2-8bit and pipenetwork/GLM-5.2-MLX-4bit, TP+RDMA or just a single Mac for the 4bit quantization).

Thank you so much @pcuenca!

[ivanfioravanti]

Same here! Thanks @pcuenca 🚀

[bibproj2]

Fixes everything for me! (EXO with 2 Mac Studios against kernelpool/GLM-5.2-8bit and pipenetwork/GLM-5.2-MLX-4bit, TP+RDMA or just a single Mac for the 4bit quantization).

Seems solved. Great work as always @pcuenca

[mitermayer]

Has the work from this PR landed in the main branch through another PR or is this issue still ongoing ?

[ivanfioravanti]

Still ongoing. Antirez reported some lower quality responses compared to GGUF and super @pcuenca was investigating.

[pcuenca]

Update:

Prompted by this thread, I replicated @antirez's quality testing protocol from DS4:

• Gather continuations for 100 prompts, from the official Z.AI GLM 5.2 API endpoint.
• Measure the negative log-likelihood of the tokens, as well as "first token" and "longest sequence" matches.
• Do the same for llama.cpp.
• Compare both.

I don't have access to @antirez's GLM code, so I checked my mxfp4 MLX quant against a Q3-K-L quant I quantized myself from the same original GLM 5.2 bf16 checkpoint, using llama.cpp.

These are the results I got (old is mlx-mxfp4, new is llamacpp-q3kl):

~/code/ds4 ❯❯❯ python compare_scores.py mlx-mxfp4.tsv llamacpp-q3kl.tsv
cases   100
tokens  2290
old_avg_nll     0.329486770
new_avg_nll     0.325934772
delta_new_minus_old     -0.003551998
relative_nll_change     -1.078%
case_wins_new_old_ties  46      54      0
first_token_matches_old_new     84      87
avg_greedy_lcp_old_new  9.310   8.620

new best cases:
case_078        delta_nll=-17.858746    tokens=24       old=0.893120    new=0.149005
case_018        delta_nll=-10.195731    tokens=24       old=0.726341    new=0.301519
case_020        delta_nll=-9.277082     tokens=24       old=0.533939    new=0.147394
case_093        delta_nll=-9.028496     tokens=24       old=0.618323    new=0.242136
case_017        delta_nll=-8.198753     tokens=24       old=0.408988    new=0.067373
case_058        delta_nll=-8.169691     tokens=24       old=0.402995    new=0.062591
case_045        delta_nll=-7.770067     tokens=24       old=0.537909    new=0.214156
case_047        delta_nll=-7.298800     tokens=24       old=0.772052    new=0.467935

old best cases:
case_088        delta_nll=14.089893     tokens=19       old=2.374772    new=3.116345
case_048        delta_nll=12.132194     tokens=24       old=0.335420    new=0.840928
case_091        delta_nll=10.788794     tokens=24       old=0.178043    new=0.627576
case_029        delta_nll=9.510132      tokens=24       old=0.227504    new=0.623759
case_005        delta_nll=9.243473      tokens=24       old=0.132364    new=0.517509
case_066        delta_nll=7.869916      tokens=24       old=0.448246    new=0.776159
case_082        delta_nll=7.583234      tokens=24       old=0.326722    new=0.642690
case_051        delta_nll=6.545374      tokens=24       old=0.294459    new=0.567183

I'm reading this as "results between Q3-K-L (llama.cpp) and mxfp4 (llama.cpp) are quite similar". At least it doesn't appear obvious to me that there are implementation errors in the MLX code (or they are subtle enough that we don't notice them in this test).

Do you have any suggestions for other quants to compare, @ivanfioravanti @antirez?

(The motivation for choosing the quants above is to be able to run them on a single M3 Ultra 512 GB computer. 4-bit llama.cpp quants OOM for me, and so does the MXFP4-MOE llama.cpp recipe. I can use distributed inference if needed, just keeping things simple for now).

[spicyneuron]

@pcuenca My mixed 4-bit MLX quant should be quite close to the Unsloth GGUF mentioned in the linked X thread.

In my quantization sweeps, over-doing some parts of the model absolutely can lead to double-digit % divergences, so I wouldn't rule that out as an explanation.

[machiabeli]

We've been maintaining #1463 in parallel (independent implementation of the same indexer_types schedule). Rather than compete, we'd like to consolidate into this PR, which has the broader tester base. Available to cherry-pick from #1463 — or we can PR directly into your branch:

• Engaged-path tests: schedule construction, per-layer cache asymmetry (CacheList(kv, indexer) on full layers vs CacheList(kv) on shared), and a test that prefills past index_topk then decodes threading prev_topk_indices layer-to-layer — i.e. the sparse path actually engaged, at B=1 and B=2.
• Strict-load validation on the real 743B mxfp4 checkpoint: 0 missing / 0 unexpected, 21 full / 57 shared.
• Independent-engine agreement: we gate a from-scratch C/Metal engine against this architecture (per-layer parity ~1e-6 at tiny config, argmax-identical full-model forward). Separately, antirez's ds4 glm5.2 branch independently implements the same 21/57 reuse-last-full-selection semantics and interleaved indexer rope — so three codebases now agree on the contested points, and the shipped checkpoint config now carries indexer_rope_interleave: true explicitly.

On the quality investigation: we're porting the same DS4 100-case protocol to our C/Metal engine, which loads the original mxfp4 weights (no re-quantization). If it scores inside the ds4 reference band, that isolates the remaining mlx-lm gap to quantization choice rather than schedule/attention implementation — and if it doesn't, that's a strong signal the gap is architectural. Happy to share results either way.

If the tests get folded in here, we'll close #1463.