Support per-layer DSA indexer schedule (GLM-5.2)

[davidrhodus]

Summary

GLM-5.2 (model_type: glm_moe_dsa) reuses the DeepSeek-V3.2 backbone but enables the DSA "lightning indexer" on only a subset of layers, declared per-layer via an indexer_types list ("full" = that layer has its own indexer weights; otherwise it has none). The shared deepseek_v32 attention builds an Indexer on every layer unconditionally, so loading GLM-5.2 fails:

ValueError: Missing N parameters:
  model.layers.*.self_attn.indexer.{wq_b,wk,weights_proj,k_norm}.weight ...

(GLM-5.2 declares "full" on layers 0,1,2 then every 4th — 22 of 78 layers.)

Changes

• deepseek_v32.py
  • DeepseekV32Attention takes layer_idx and builds the Indexer only when indexer_types[layer_idx] == "full" — or when indexer_types is absent, in which case the indexer is built on every layer (DeepSeek-V3.2 behaviour is unchanged).
  • forward guards the indexer/sparse-mask block behind self.indexer is not None.
  • make_cache gives indexer-less layers a single-slot CacheList(KVCache()) — a second, never-written KVCache otherwise crashes CacheList.state on its empty keys.
• glm_moe_dsa.py — add the optional indexer_types config field.
• tests/test_models.py — glm_moe_dsa config with mixed ["full","shared",...] indexer_types, exercising both the conditional indexer build and the mixed-cache decode path.

Testing

• New test passes the standard model_test_runner checks (fp32/fp16 forward, KV-cache decode, batch > 1, deepcopy); confirmed indexer is built only on "full" layers and the single-slot cache decodes correctly.
• End-to-end: converted the full GLM-5.2 (glm_moe_dsa, 256 experts) to MLX and verified coherent generation (it failed with "Missing parameters" before this change).

DeepSeek-V3.2 / other glm_moe_dsa checkpoints without indexer_types are unaffected (indexer still built on every layer, 2-slot cache).

🤖 Generated with Claude Code