atlas_core/config/parsers/
glm5_next.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! GLM-5.3-Flash (`glm5_next`) config parser.
4//!
5//! Reference checkpoint: `LibertAIDAI/GLM-5.3-Flash-NVFP4` snapshot
6//! `9e0d74e3cef17f634e84fb8e2223707e02616290`. All constants asserted in the
7//! tests below were read from that checkpoint's `config.json` and from a scan
8//! of all 120 shard headers. Claims are scoped to that checkpoint.
9//!
10//! Two things make this model unlike anything already parsed here:
11//!
12//! 1. **NoPE MLA.** `qk_rope_head_dim` is `0`. The DeepSeek-V4 parser derives
13//!    `qk_nope_head_dim = head_dim - qk_rope_head_dim` and only does so when
14//!    `qk_rope_head_dim > 0`, i.e. it encodes "there is a rope section" as an
15//!    invariant. We must NOT reuse it: a zero here is a fact to preserve, not a
16//!    missing value to repair. Upstream vLLM's equivalent assumption is exactly
17//!    what made its SM120 sparse-MLA backend structurally unusable for this
18//!    model (`pe_dim must be 64 for fp8_ds_mla`).
19//!
20//! 2. **Hybrid layer stack.** 45 text layers alternate KDA linear attention
21//!    with DeepSeek-style sparse attention, and the checkpoint states the split
22//!    explicitly in `linear_attn_config.{kda_layers,full_attn_layers}`. We trust
23//!    that list over any modular arithmetic, and cross-check it against the
24//!    `layer_types` array when present.
25
26use anyhow::{Context, Result, bail};
27
28use super::super::{Glm5NextRouterMode, LayerType, ModelConfig, finalize_config};
29
30/// `num_hidden_layers` counts text layers only; the MTP layer sits at index
31/// `num_hidden_layers` (45) and is NOT included in that count.
32///
33/// 🪤 GLM-5.3 does **not** use DeepSeek's `mtp.0.*` naming — the MTP weights
34/// live under `model.language_model.layers.45.*`. A `grep mtp` over this
35/// checkpoint returns zero hits. Verified over all 120 shard headers.
36pub fn glm5_next_mtp_layer_index(config: &ModelConfig) -> usize {
37    config.num_hidden_layers
38}
39
40fn text_config(raw: &serde_json::Value) -> &serde_json::Value {
41    raw.get("text_config").unwrap_or(raw)
42}
43
44/// GLM's own name for its sparse-MLA mixer.
45///
46/// Since Slice 8 this maps to [`LayerType::SparseAttention`], a real variant, and the
47/// array **round-trips**: `layer_types[i].hf_name()` reproduces the checkpoint string.
48/// It used to be flattened onto `FullAttention` "for scheduling purposes" — which was
49/// only true while nothing scheduled on it. A sparse layer needs indexer state, an
50/// indexer weight family and a per-query top-k step, so cache sizing and weight binding
51/// have to be able to tell the two apart.
52pub const GLM5NEXT_SPARSE_ATTN: &str = "deepseek_sparse_attention";
53
54pub fn parse_glm5_next(json: &str) -> Result<ModelConfig> {
55    let raw: serde_json::Value =
56        serde_json::from_str(json).context("Invalid JSON in GLM-5.3 (glm5_next) config.json")?;
57
58    // GLM-5.3 nests everything under `text_config` (the top level carries the
59    // multimodal wrapper + `quantization_config`). Parse the inner object, but
60    // keep the outer one for quantization and for the architecture string.
61    let text = text_config(&raw).clone();
62
63    // `deepseek_sparse_attention` now deserializes onto `LayerType::SparseAttention` via a
64    // serde alias, so the strip below is no longer required for parsing to succeed. It stays
65    // because `build_layer_types` derives the array from `linear_attn_config`'s index lists
66    // (the authoritative source) and then CROSS-CHECKS it against the textual array; letting
67    // serde populate the field first would make that cross-check compare a value against
68    // itself.
69    let mut text_for_struct = text.clone();
70    if let Some(obj) = text_for_struct.as_object_mut() {
71        obj.remove("layer_types");
72        // 🔴 GLM-5.3-Flash declares THREE stop tokens — `eos_token_id` is an ARRAY:
73        // 154820 `<|endoftext|>`, 154827 `<|user|>`, 154829 `<|observation|>`. Before Slice 9
74        // that made this checkpoint's config.json fail to parse outright ("invalid type:
75        // sequence, expected u32"); the parser had only ever seen a hand-written fixture with
76        // no `eos_token_id` at all. Handled centrally now — `eos_token_id_field` takes element
77        // 0 as the primary (`tokenizer_config.json` names `<|endoftext|>` as THE `eos_token`)
78        // and `parse_config` fills `eos_token_ids` with the complete set.
79        // Read stop tokens through `ModelConfig::eos_ids()` / `is_eos()`, never off the scalar.
80    }
81    let text_json =
82        serde_json::to_string(&text_for_struct).context("re-serialize glm5_next text_config")?;
83    let mut config: ModelConfig =
84        serde_json::from_str(&text_json).context("Failed to parse glm5_next text_config")?;
85    let text = &text;
86
87    // Canonical model_type. The inner object says `glm5_next_text`; Atlas keys
88    // dispatch off the outer family name.
89    config.model_type = "glm5_next".to_string();
90
91    // ---- MLA geometry -----------------------------------------------------
92    // NoPE: qk_rope_head_dim is legitimately 0 and must survive untouched.
93    // Do not "repair" it, and do not derive qk_nope from it.
94    let qk_rope = text
95        .get("qk_rope_head_dim")
96        .and_then(|v| v.as_u64())
97        .map(|v| v as usize);
98    match qk_rope {
99        Some(v) => config.qk_rope_head_dim = v,
100        // Absent (not zero) would mean an unfamiliar variant: refuse rather
101        // than silently assume NoPE.
102        None => bail!(
103            "glm5_next config.json has no qk_rope_head_dim; refusing to guess \
104             whether this checkpoint is NoPE"
105        ),
106    }
107    if let Some(v) = text.get("qk_nope_head_dim").and_then(|v| v.as_u64()) {
108        config.qk_nope_head_dim = v as usize;
109    }
110    if let Some(v) = text.get("v_head_dim").and_then(|v| v.as_u64()) {
111        config.v_head_dim = v as usize;
112    }
113    // `head_dim` is 0 in the checkpoint. For MLA the meaningful per-head width
114    // is qk_head_dim (256), NOT hidden_size / num_attention_heads (which would
115    // give 64 and silently corrupt every attention shape) — the same trap the
116    // DeepSeek-V4 parser documents for its own head_dim.
117    if config.head_dim == 0 {
118        config.head_dim = text
119            .get("qk_head_dim")
120            .and_then(|v| v.as_u64())
121            .map(|v| v as usize)
122            .unwrap_or(config.qk_nope_head_dim + config.qk_rope_head_dim);
123    }
124    // partial_rotary_factor is rope_dim / head_dim; NoPE makes it exactly 0.
125    config.partial_rotary_factor = if config.head_dim > 0 {
126        config.qk_rope_head_dim as f64 / config.head_dim as f64
127    } else {
128        0.0
129    };
130
131    // ---- MoE --------------------------------------------------------------
132    if config.num_experts == 0 && config.n_routed_experts > 0 {
133        config.num_experts = config.n_routed_experts;
134    }
135    let n_shared = text
136        .get("n_shared_experts")
137        .and_then(|v| v.as_u64())
138        .unwrap_or(0) as usize;
139    if config.shared_expert_intermediate_size == 0 && n_shared > 0 {
140        config.shared_expert_intermediate_size = n_shared * config.moe_intermediate_size;
141    }
142
143    // ---- KDA linear attention --------------------------------------------
144    // The checkpoint carries these under `linear_attn_config`, not as the
145    // flat `linear_*` keys Atlas uses for Qwen GDN.
146    let lac = text.get("linear_attn_config");
147    if let Some(lac) = lac {
148        let g = |k: &str| lac.get(k).and_then(|v| v.as_u64()).map(|v| v as usize);
149        if let Some(v) = g("num_heads") {
150            config.linear_num_key_heads = v;
151            config.linear_num_value_heads = v;
152        }
153        if let Some(v) = g("head_dim") {
154            config.linear_key_head_dim = v;
155            config.linear_value_head_dim = v;
156        }
157        if let Some(v) = g("short_conv_kernel_size") {
158            config.linear_conv_kernel_dim = v;
159        }
160        // Bounds the log-decay `kda_gate` emits. Absent means an unfamiliar KDA variant, not
161        // "no bound" — 0.0 would clamp every decay to <= 0 and is a different model.
162        match lac.get("gate_lower_bound").and_then(|v| v.as_f64()) {
163            Some(v) => config.linear_gate_lower_bound = v as f32,
164            None => bail!(
165                "glm5_next: linear_attn_config has no gate_lower_bound; refusing to guess the \
166                 KDA decay bound (GLM-5.3-Flash declares -5.0)"
167            ),
168        }
169    }
170
171    // ---- DSA indexer ------------------------------------------------------
172    // Default only if truly absent. GLM's value is 2048 and it matters: it is
173    // part of what a sparse-MLA backend gates on.
174    if config.index_topk == 0
175        && let Some(v) = text.get("index_topk").and_then(|v| v.as_u64())
176    {
177        config.index_topk = v as usize;
178    }
179    // `index_kpool` sets the pool budget (`index_topk / index_kpool`), so a missing
180    // value is not a cosmetic default — it would silently change how many candidates
181    // the top-k ranks. Read it, never guess it.
182    if let Some(v) = text.get("index_kpool").and_then(|v| v.as_u64()) {
183        config.index_kpool = v as usize;
184    }
185    if let Some(v) = text
186        .get("index_kpool_always_select_tail")
187        .and_then(|v| v.as_bool())
188    {
189        config.index_kpool_always_select_tail = v;
190    }
191
192    // ---- Layer types ------------------------------------------------------
193    config.layer_types = build_layer_types(text, config.num_hidden_layers)?;
194
195    // MTP / NextN layers sit PAST the text stack. GLM-5.3-Flash declares
196    // `num_nextn_predict_layers = 1`, so layer 45 exists as a real decoder layer while
197    // `layer_types` legitimately covers only 0..=44. Give it its own slot rather than
198    // appending it, so every "iterate the text stack" loop keeps meaning what it says.
199    //
200    // The config does not state the MTP block's mixer kind. Derived here as "the same
201    // non-linear mixer this model uses", which the Slice-8 checkpoint audit confirms:
202    // layer 45's `self_attn` tensor set is name/dtype/shape IDENTICAL to the 11
203    // `deepseek_sparse_attention` text layers. Weight binding classifies from tensor
204    // names anyway and does not trust this field.
205    let n_mtp = text
206        .get("num_nextn_predict_layers")
207        .and_then(|v| v.as_u64())
208        .unwrap_or(0) as usize;
209    if n_mtp > 0 {
210        let kind = if config.layer_types.contains(&LayerType::SparseAttention) {
211            LayerType::SparseAttention
212        } else {
213            LayerType::FullAttention
214        };
215        config.mtp_layer_types = vec![kind; n_mtp];
216    }
217
218    // ---- Dense-vs-routed MLP split ---------------------------------------
219    // `first_k_dense_replace = 3`: layers 0..=2 carry a dense MLP, every later text layer
220    // routes to experts. Nothing in Atlas read this before, so `mlp_only_layers` came out
221    // EMPTY and the whole stack looked routed — a dense layer bound as MoE looks for
222    // `mlp.experts.*` that do not exist. Cross-checked against the textual
223    // `mlp_layer_types` array when the checkpoint carries one, the same way
224    // `build_layer_types` cross-checks the mixer map.
225    config.mlp_only_layers = build_mlp_only_layers(text, config.num_hidden_layers)?;
226
227    // ---- SwiGLU clamp -----------------------------------------------------
228    // 🔴 GLM clamps its SwiGLU and the clamp is ASYMMETRIC (`gate` upper-bounded only, `up`
229    // both ways). Nothing in Atlas read this before, so the value would have had to be
230    // hardcoded at a call site or defaulted to "no clamp" — and an absent clamp is INVISIBLE
231    // on well-scaled activations, firing only on the tails. Read it, never guess it.
232    config.swiglu_limit = match text.get("swiglu_limit").and_then(|v| v.as_f64()) {
233        Some(v) if v > 0.0 => v as f32,
234        Some(v) => bail!("glm5_next: swiglu_limit is {v}, which cannot bound anything"),
235        None => bail!(
236            "glm5_next config.json has no swiglu_limit; refusing to guess whether this \
237             checkpoint clamps its SwiGLU (GLM-5.3-Flash declares 10.0)"
238        ),
239    };
240
241    // ---- Grouped expert routing -------------------------------------------
242    // `n_group`/`topk_group` are 1 here, which makes the group mask all-ones and grouped
243    // routing a NO-OP. `glm5next_router_topk` implements plain top-k and takes `n_group`
244    // only to REFUSE a checkpoint where the machinery would matter. Refuse at parse time
245    // too, so the failure names the config rather than surfacing as a silent kernel return.
246    for key in ["n_group", "topk_group"] {
247        if let Some(v) = text.get(key).and_then(|v| v.as_u64())
248            && v != 1
249        {
250            bail!(
251                "glm5_next: {key} = {v}. Grouped expert routing is not implemented — \
252                 glm5next_router_topk ranks every expert in one group."
253            );
254        }
255    }
256
257    // ---- Router dtype ladder ---------------------------------------------
258    // 🔴 SEMANTIC, not precision. Default = HF 5.16.1's explicit fp32 router. The checkpoint may
259    // override with `moe_router_dtype` (the same field vLLM reads); an UNRECOGNISED value is a
260    // hard error, never a silent fallback onto the other semantics.
261    config.glm5next_router_mode = match text.get("moe_router_dtype").and_then(|v| v.as_str()) {
262        None => Glm5NextRouterMode::HfFp32,
263        Some(s) => Glm5NextRouterMode::from_config_str(s).ok_or_else(|| {
264            anyhow::anyhow!(
265                "glm5_next: unknown moe_router_dtype {s:?}; expected float32 or bfloat16"
266            )
267        })?,
268    };
269
270    finalize_config(&mut config, &raw).context("glm5_next: finalize_config")?;
271    refuse_shared_indexer(text, &config).context("glm5_next: indexer_types")?;
272    validate_glm5_next(&config)?;
273    Ok(config)
274}
275
276/// Layers whose MLP is dense rather than routed.
277///
278/// `first_k_dense_replace` is the authoritative knob; the textual `mlp_layer_types` array is
279/// used to CROSS-CHECK it, never as a silent substitute. A disagreement is a hard error: the
280/// two answers differing means the checkpoint is not the one this parser was written for.
281fn build_mlp_only_layers(text: &serde_json::Value, n_layers: usize) -> Result<Vec<usize>> {
282    let first_k = text
283        .get("first_k_dense_replace")
284        .and_then(|v| v.as_u64())
285        .map(|v| v as usize);
286    let textual: Option<Vec<usize>> =
287        text.get("mlp_layer_types")
288            .and_then(|v| v.as_array())
289            .map(|a| {
290                a.iter()
291                    .enumerate()
292                    .filter(|(_, v)| v.as_str() != Some("sparse"))
293                    .map(|(i, _)| i)
294                    .collect()
295            });
296
297    match (first_k, textual) {
298        (Some(k), t) => {
299            if k > n_layers {
300                bail!("glm5_next: first_k_dense_replace {k} exceeds num_hidden_layers {n_layers}");
301            }
302            let derived: Vec<usize> = (0..k).collect();
303            if let Some(t) = t
304                && t != derived
305            {
306                bail!(
307                    "glm5_next: first_k_dense_replace={k} implies dense layers \
308                     {derived:?}, but mlp_layer_types says {t:?}"
309                );
310            }
311            Ok(derived)
312        }
313        // No `first_k_dense_replace`: the textual array is then the only statement of the
314        // split, and it must be present and correctly sized.
315        (None, Some(t)) => Ok(t),
316        (None, None) => bail!(
317            "glm5_next: neither first_k_dense_replace nor mlp_layer_types present; \
318             refusing to guess which layers are dense"
319        ),
320    }
321}
322
323/// Build the per-layer mixer map.
324///
325/// Priority: the explicit `linear_attn_config.kda_layers` / `full_attn_layers`
326/// lists, cross-checked against `layer_types` when both are present. We do not
327/// fall back to `layer % 4 == 3` arithmetic — that pattern happens to hold for
328/// this checkpoint but is not stated anywhere as a contract.
329fn build_layer_types(text: &serde_json::Value, n_layers: usize) -> Result<Vec<LayerType>> {
330    let idx_list = |key: &str| -> Option<Vec<usize>> {
331        text.get("linear_attn_config")?
332            .get(key)?
333            .as_array()
334            .map(|a| {
335                a.iter()
336                    .filter_map(|v| v.as_u64())
337                    .map(|v| v as usize)
338                    .collect()
339            })
340    };
341
342    let kda = idx_list("kda_layers");
343    let full = idx_list("full_attn_layers");
344
345    let mut types = vec![LayerType::FullAttention; n_layers];
346    match (kda, full) {
347        (Some(kda), Some(full)) => {
348            if kda.len() + full.len() != n_layers {
349                bail!(
350                    "glm5_next: kda_layers ({}) + full_attn_layers ({}) != num_hidden_layers ({})",
351                    kda.len(),
352                    full.len(),
353                    n_layers
354                );
355            }
356            for i in &kda {
357                if *i >= n_layers {
358                    bail!("glm5_next: kda_layers index {i} out of range for {n_layers} layers");
359                }
360                types[*i] = LayerType::LinearAttention;
361            }
362            // `full_attn_layers` is GLM's name for "not KDA". On GLM-5.3-Flash those layers
363            // are `deepseek_sparse_attention`, so the textual array decides the variant —
364            // the index list alone cannot tell sparse from dense full attention.
365            let textual = text.get("layer_types").and_then(|v| v.as_array());
366            for i in &full {
367                if *i >= n_layers {
368                    bail!("glm5_next: full_attn_layers index {i} out of range");
369                }
370                if types[*i] == LayerType::LinearAttention {
371                    bail!("glm5_next: layer {i} listed as BOTH kda and full attention");
372                }
373                types[*i] = match textual.and_then(|a| a.get(*i)).and_then(|v| v.as_str()) {
374                    Some(GLM5NEXT_SPARSE_ATTN) => LayerType::SparseAttention,
375                    _ => LayerType::FullAttention,
376                };
377            }
378        }
379        _ => {
380            // Fall back to the textual `layer_types` array.
381            let arr = text
382                .get("layer_types")
383                .and_then(|v| v.as_array())
384                .context("glm5_next: neither linear_attn_config lists nor layer_types present")?;
385            if arr.len() != n_layers {
386                bail!(
387                    "glm5_next: layer_types has {} entries, expected {n_layers}",
388                    arr.len()
389                );
390            }
391            for (i, v) in arr.iter().enumerate() {
392                types[i] = match v.as_str().unwrap_or("") {
393                    "linear_attention" => LayerType::LinearAttention,
394                    GLM5NEXT_SPARSE_ATTN => LayerType::SparseAttention,
395                    "full_attention" => LayerType::FullAttention,
396                    other => bail!("glm5_next: unknown layer_type {other:?} at layer {i}"),
397                };
398            }
399        }
400    }
401
402    // Cross-check against layer_types when we used the index lists.
403    if let Some(arr) = text.get("layer_types").and_then(|v| v.as_array())
404        && arr.len() == n_layers
405    {
406        for (i, v) in arr.iter().enumerate() {
407            let want = match v.as_str().unwrap_or("") {
408                "linear_attention" => LayerType::LinearAttention,
409                GLM5NEXT_SPARSE_ATTN => LayerType::SparseAttention,
410                _ => LayerType::FullAttention,
411            };
412            if types[i] != want {
413                bail!(
414                    "glm5_next: layer {i} disagrees — index lists say {:?}, layer_types says {want:?}",
415                    types[i]
416                );
417            }
418        }
419    }
420    Ok(types)
421}
422
423/// Refuse a checkpoint whose DSA layers use **shared** indexing.
424///
425/// `indexer_types[i]` is `"full"` (this layer runs its own DSA indexer) or `"shared"`
426/// (it reuses the previous full layer's top-k selection, which HF propagates as
427/// `prev_topk_indices`). A shared layer must attend to the **upstream** layer's token
428/// set; running its own indexer instead yields a plausible wrong answer with no crash
429/// and no shape error, so this is refused at load rather than discovered in output.
430///
431/// `LibertAIDAI/GLM-5.3-Flash-NVFP4@9e0d74e3` carries an explicit 45-entry array that is
432/// entirely `"full"` — its 11 DSA layers (3, 7, … 43) each run their own indexer — so
433/// propagation is deliberately NOT implemented. Verified against the checkpoint config
434/// 2026-08-27.
435///
436/// Nothing is stored: the value of this check is the refusal. Adding a config field no
437/// runtime path reads would be dead surface.
438///
439/// When the array is absent, HF derives it, and so must we — an absent array does not
440/// mean all-full. Source: `transformers` 5.16.1 `configuration_glm5_next.py`:
441/// `index_topk_pattern` (an `"FSSF…"` string) wins, else
442/// `"full" if (max(i - offset + 1, 0) % freq) == 0` with `freq = max(index_topk_freq, 1)`
443/// and `offset = index_skip_topk_offset` (default 2).
444fn refuse_shared_indexer(text: &serde_json::Value, config: &ModelConfig) -> Result<()> {
445    let n = config.num_hidden_layers;
446    let modes: Vec<String> = if let Some(arr) = text.get("indexer_types").and_then(|v| v.as_array())
447    {
448        if arr.len() != n {
449            bail!(
450                "indexer_types has {} entries for {n} layers; a length mismatch would \
451                 silently misalign every layer's indexer mode",
452                arr.len()
453            );
454        }
455        arr.iter()
456            .map(|v| {
457                v.as_str()
458                    .map(str::to_string)
459                    .context("indexer_types entry is not a string")
460            })
461            .collect::<Result<_>>()?
462    } else if let Some(pat) = text.get("index_topk_pattern").and_then(|v| v.as_str()) {
463        if pat.chars().count() != n {
464            bail!(
465                "index_topk_pattern has {} chars for {n} layers",
466                pat.chars().count()
467            );
468        }
469        pat.chars()
470            .map(|c| match c {
471                'F' => Ok("full".to_string()),
472                'S' => Ok("shared".to_string()),
473                other => bail!("index_topk_pattern: unknown char {other:?}, expected F or S"),
474            })
475            .collect::<Result<_>>()?
476    } else {
477        let freq = text
478            .get("index_topk_freq")
479            .and_then(|v| v.as_u64())
480            .unwrap_or(1)
481            .max(1) as i64;
482        let offset = text
483            .get("index_skip_topk_offset")
484            .and_then(|v| v.as_i64())
485            .unwrap_or(2);
486        (0..n)
487            .map(|i| {
488                let shifted = ((i as i64) - offset + 1).max(0);
489                if shifted % freq == 0 {
490                    "full"
491                } else {
492                    "shared"
493                }
494                .to_string()
495            })
496            .collect()
497    };
498
499    for (i, m) in modes.iter().enumerate() {
500        if m != "full" && m != "shared" {
501            bail!("layer {i}: unknown indexer mode {m:?}, expected \"full\" or \"shared\"");
502        }
503    }
504
505    // Only DSA layers have an indexer at all; a "shared" entry on a KDA layer is inert.
506    let shared: Vec<usize> = config
507        .layer_types
508        .iter()
509        .enumerate()
510        .filter(|(i, t)| {
511            **t != LayerType::LinearAttention && modes.get(*i).is_some_and(|m| m == "shared")
512        })
513        .map(|(i, _)| i)
514        .collect();
515    if !shared.is_empty() {
516        bail!(
517            "DSA layer(s) {shared:?} use SHARED indexing (reuse the previous full layer's \
518             top-k). Atlas runs a per-layer indexer and does not propagate selections, so \
519             these layers would attend to the wrong token set — a wrong answer, not a \
520             crash. GLM-5.3-Flash-NVFP4 is entirely \"full\"; implement prev_topk_indices \
521             propagation before serving a checkpoint that is not."
522        );
523    }
524    Ok(())
525}
526
527fn validate_glm5_next(config: &ModelConfig) -> Result<()> {
528    if config.qk_rope_head_dim != 0 {
529        bail!(
530            "glm5_next: expected NoPE (qk_rope_head_dim == 0), got {}. \
531             A non-zero value means this is not the GLM-5.3 geometry we support.",
532            config.qk_rope_head_dim
533        );
534    }
535    if config.head_dim == 0 {
536        bail!("glm5_next: head_dim resolved to 0");
537    }
538    let linear = config
539        .layer_types
540        .iter()
541        .filter(|t| **t == LayerType::LinearAttention)
542        .count();
543    let full = config.layer_types.len() - linear;
544    if linear == 0 || full == 0 {
545        bail!("glm5_next: degenerate layer map — {linear} linear / {full} full");
546    }
547    Ok(())
548}
549
550#[cfg(test)]
551mod tests {
552    use super::*;
553
554    /// Trimmed from the real checkpoint config.json (LibertAI NVFP4 @ 9e0d74e3).
555    /// Values are verbatim; only unrelated keys were dropped for size.
556    fn glm53_config_json() -> String {
557        // KDA layers are every index NOT congruent to 3 (mod 4); the real file
558        // enumerates them explicitly and we mirror that here.
559        let kda: Vec<String> = (0..45)
560            .filter(|i| i % 4 != 3)
561            .map(|i| i.to_string())
562            .collect();
563        let full: Vec<String> = (0..45)
564            .filter(|i| i % 4 == 3)
565            .map(|i| i.to_string())
566            .collect();
567        let layer_types: Vec<String> = (0..45)
568            .map(|i| {
569                if i % 4 == 3 {
570                    "\"deepseek_sparse_attention\"".to_string()
571                } else {
572                    "\"linear_attention\"".to_string()
573                }
574            })
575            .collect();
576        format!(
577            r#"{{
578  "architectures": ["Glm5NextForConditionalGeneration"],
579  "model_type": "glm5_next",
580  "text_config": {{
581    "model_type": "glm5_next_text",
582    "num_hidden_layers": 45,
583    "num_nextn_predict_layers": 1,
584    "hidden_size": 4096,
585    "intermediate_size": 12288,
586    "num_attention_heads": 64,
587    "num_key_value_heads": 64,
588    "head_dim": 0,
589    "qk_head_dim": 256,
590    "qk_nope_head_dim": 256,
591    "qk_rope_head_dim": 0,
592    "v_head_dim": 256,
593    "kv_lora_rank": 512,
594    "q_lora_rank": 1536,
595    "mla_use_nope": true,
596    "index_topk": 2048,
597    "index_kpool": 4,
598    "index_n_heads": 32,
599    "index_head_dim": 128,
600    "hc_mult": 4,
601    "hc_sinkhorn_iters": 20,
602    "hc_eps": 1e-06,
603    "mhc": true,
604    "n_routed_experts": 288,
605    "n_shared_experts": 1,
606    "num_experts_per_tok": 8,
607    "moe_intermediate_size": 2048,
608    "first_k_dense_replace": 3,
609    "swiglu_limit": 10.0,
610    "routed_scaling_factor": 2.5,
611    "norm_topk_prob": true,
612    "n_group": 1,
613    "topk_group": 1,
614    "scoring_func": "sigmoid",
615    "topk_method": "noaux_tc",
616    "rms_norm_eps": 1e-05,
617    "vocab_size": 154880,
618    "max_position_embeddings": 1048576,
619    "linear_attn_config": {{
620      "num_heads": 64,
621      "head_dim": 128,
622      "short_conv_kernel_size": 4,
623      "gate_lower_bound": -5.0,
624      "kda_layers": [{kda}],
625      "full_attn_layers": [{full}]
626    }},
627    "layer_types": [{lt}]
628  }}
629}}"#,
630            kda = kda.join(","),
631            full = full.join(","),
632            lt = layer_types.join(",")
633        )
634    }
635
636    /// Inject a key into `text_config` and re-serialise, so a test can vary one
637    /// checkpoint field without restating the whole fixture.
638    fn with_text_key(key: &str, value: serde_json::Value) -> String {
639        let mut raw: serde_json::Value =
640            serde_json::from_str(&glm53_config_json()).expect("fixture json");
641        raw["text_config"][key] = value;
642        raw.to_string()
643    }
644
645    /// The real checkpoint's array: 45 entries, every one `"full"`.
646    #[test]
647    fn an_all_full_indexer_array_is_accepted() {
648        let all_full = serde_json::Value::from(vec!["full"; 45]);
649        let c = parse_glm5_next(&with_text_key("indexer_types", all_full)).expect("parse");
650        // The 11 DSA layers of GLM-5.3-Flash.
651        let dsa: Vec<usize> = c
652            .layer_types
653            .iter()
654            .enumerate()
655            .filter(|(_, t)| **t != LayerType::LinearAttention)
656            .map(|(i, _)| i)
657            .collect();
658        assert_eq!(dsa, vec![3, 7, 11, 15, 19, 23, 27, 31, 35, 39, 43]);
659    }
660
661    /// 🔴 A shared DSA layer must attend to the UPSTREAM layer's token set. Atlas runs a
662    /// per-layer indexer and does not propagate, so this is a wrong answer with no crash
663    /// — refused at load.
664    #[test]
665    fn a_shared_dsa_layer_is_refused() {
666        let mut modes = vec!["full"; 45];
667        modes[7] = "shared"; // a real DSA layer
668        let e = parse_glm5_next(&with_text_key("indexer_types", modes.into()))
669            .expect_err("shared DSA indexing must be refused");
670        let msg = e.to_string() + &e.root_cause().to_string();
671        assert!(msg.contains('7'), "the error must name the layer: {msg}");
672    }
673
674    /// A `"shared"` entry on a KDA layer is inert — those layers have no indexer at all.
675    /// Refusing it would reject a legal checkpoint.
676    #[test]
677    fn a_shared_entry_on_a_linear_layer_is_inert() {
678        let mut modes = vec!["full"; 45];
679        modes[0] = "shared"; // layer 0 is KDA
680        assert!(parse_glm5_next(&with_text_key("indexer_types", modes.into())).is_ok());
681    }
682
683    /// A length mismatch would silently misalign every layer's mode.
684    #[test]
685    fn a_wrong_length_indexer_array_is_refused() {
686        let short = serde_json::Value::from(vec!["full"; 44]);
687        assert!(parse_glm5_next(&with_text_key("indexer_types", short)).is_err());
688    }
689
690    /// An ABSENT array does not mean all-full: HF derives it. With the default
691    /// `freq = 1` every layer is full, which is why the bare fixture parses — but a
692    /// `freq = 4` schedule genuinely produces shared DSA layers and must be refused.
693    #[test]
694    fn an_absent_array_is_derived_not_assumed_full() {
695        // freq=1 (the default) → all full → parses.
696        assert!(parse_glm5_next(&glm53_config_json()).is_ok());
697        // freq=4, offset=2 → full only where (max(i-1,0) % 4 == 0); DSA layer 7 is shared.
698        let e = parse_glm5_next(&with_text_key("index_topk_freq", 4.into()))
699            .expect_err("a freq schedule that shares DSA layers must be refused");
700        assert!(e.root_cause().to_string().contains("SHARED"), "{e}");
701    }
702
703    /// The `"FSSF…"` pattern string overrides the freq schedule, as in HF.
704    #[test]
705    fn an_index_topk_pattern_is_honoured() {
706        // All-F pattern parses; one S on a DSA layer does not.
707        let ok: String = "F".repeat(45);
708        assert!(parse_glm5_next(&with_text_key("index_topk_pattern", ok.into())).is_ok());
709        let mut bad: Vec<char> = "F".repeat(45).chars().collect();
710        bad[43] = 'S'; // the last DSA layer
711        let bad: String = bad.into_iter().collect();
712        assert!(parse_glm5_next(&with_text_key("index_topk_pattern", bad.into())).is_err());
713    }
714
715    #[test]
716    fn parses_glm5_next() {
717        let c = parse_glm5_next(&glm53_config_json()).expect("parse");
718        assert_eq!(c.model_type, "glm5_next");
719        assert_eq!(c.num_hidden_layers, 45);
720        assert_eq!(c.hidden_size, 4096);
721        assert_eq!(c.n_routed_experts, 288);
722        assert_eq!(c.num_experts_per_tok, 8);
723        assert_eq!(c.moe_intermediate_size, 2048);
724    }
725
726    /// The acceptance criterion: a legitimate zero must round-trip untouched.
727    #[test]
728    fn nope_rope_dim_zero_survives_exactly() {
729        let c = parse_glm5_next(&glm53_config_json()).expect("parse");
730        assert_eq!(c.qk_rope_head_dim, 0, "NoPE zero must not be 'repaired'");
731        assert_eq!(c.qk_nope_head_dim, 256, "nope dim must come from the file");
732        assert_eq!(c.v_head_dim, 256);
733        assert_eq!(c.partial_rotary_factor, 0.0);
734    }
735
736    /// 🪤 GLM-5.3 is **MLA AND NoPE at the same time** — the exact combination
737    /// that broke Atlas's MLA decode dispatch.
738    ///
739    /// `qwen3_attention/decode/run_paged_decode.rs` used to select the MLA
740    /// compressed-cache decode with `mla.rope > 0`, treating "has a RoPE
741    /// section" as a proxy for "is MLA". That holds for DeepSeek-V4-Flash
742    /// (rope=64) and fails here: GLM has a 512-dim latent cache and rope=0, so
743    /// it satisfied "is MLA" while failing the proxy, and would have fallen
744    /// through to the generic *non-MLA* paged decode carrying an MLA-shaped
745    /// cache — a wrong answer with no crash.
746    ///
747    /// The predicate is now `mla.is_some()`, and only the genuine RoPE
748    /// operations are guarded on `rope > 0`. This test pins the config-level
749    /// fact that makes the distinction necessary; if it ever fails, re-read that
750    /// dispatch before touching anything else.
751    #[test]
752    fn is_mla_and_nope_simultaneously() {
753        let c = parse_glm5_next(&glm53_config_json()).expect("parse");
754        assert!(
755            c.kv_lora_rank > 0,
756            "GLM-5.3 is MLA: a latent KV cache of {} dims",
757            c.kv_lora_rank
758        );
759        assert_eq!(c.kv_lora_rank, 512);
760        assert_eq!(
761            c.qk_rope_head_dim, 0,
762            "...and simultaneously NoPE. `rope > 0` must never stand in for `is MLA`."
763        );
764    }
765
766    /// head_dim=0 in-file must resolve to qk_head_dim (256), never to
767    /// hidden_size / num_attention_heads (64).
768    #[test]
769    fn head_dim_resolves_to_mla_width_not_hidden_over_heads() {
770        let c = parse_glm5_next(&glm53_config_json()).expect("parse");
771        assert_eq!(c.head_dim, 256);
772        assert_ne!(c.head_dim, 4096 / 64);
773    }
774
775    /// Reconciled census, TEXT LAYERS ONLY (0..44): 34 KDA + 11 DSA.
776    /// See .planning/ATLAS-GLM5NEXT-SKILL-RECONCILIATION-20260826.md — the
777    /// whole-checkpoint totals differ because layer 45 (MTP) is DSA-shaped.
778    #[test]
779    fn layer_census_matches_reconciled_counts() {
780        let c = parse_glm5_next(&glm53_config_json()).expect("parse");
781        let kda = c
782            .layer_types
783            .iter()
784            .filter(|t| **t == LayerType::LinearAttention)
785            .count();
786        // Since Slice 8 the DSA layers are `SparseAttention`, not `FullAttention` — and
787        // there must be ZERO plain full-attention layers, or something was flattened.
788        let dsa = c
789            .layer_types
790            .iter()
791            .filter(|t| **t == LayerType::SparseAttention)
792            .count();
793        let plain_full = c
794            .layer_types
795            .iter()
796            .filter(|t| **t == LayerType::FullAttention)
797            .count();
798        assert_eq!(c.layer_types.len(), 45);
799        assert_eq!(kda, 34, "KDA layers over text layers 0..44");
800        assert_eq!(dsa, 11, "DSA layers over text layers 0..44");
801        assert_eq!(plain_full, 0, "GLM-5.3 has no plain full-attention layer");
802        // Spot-check the actual indices, not just the totals.
803        assert_eq!(c.layer_types[0], LayerType::LinearAttention);
804        assert_eq!(c.layer_types[3], LayerType::SparseAttention);
805        assert_eq!(c.layer_types[43], LayerType::SparseAttention);
806        assert_eq!(c.layer_types[44], LayerType::LinearAttention);
807    }
808
809    /// `layer_types` must round-trip back to the checkpoint's own vocabulary. This is what
810    /// "flattened onto FullAttention" used to break: the parse succeeded and the array
811    /// silently said `full_attention` where the checkpoint said `deepseek_sparse_attention`.
812    #[test]
813    fn layer_types_round_trip_to_the_checkpoint_strings() {
814        let c = parse_glm5_next(&glm53_config_json()).expect("parse");
815        let raw: serde_json::Value = serde_json::from_str(&glm53_config_json()).unwrap();
816        let want = raw["text_config"]["layer_types"]
817            .as_array()
818            .expect("layer_types");
819        assert_eq!(want.len(), c.layer_types.len());
820        for (i, w) in want.iter().enumerate() {
821            assert_eq!(
822                c.layer_types[i].hf_name(),
823                w.as_str().unwrap(),
824                "layer {i} does not round-trip"
825            );
826        }
827    }
828
829    /// Layer 45 (MTP) is a real decoder layer that is NOT part of the text stack.
830    /// It must be representable without being appended to `layer_types`.
831    #[test]
832    fn mtp_layer_is_represented_outside_the_text_stack() {
833        let c = parse_glm5_next(&glm53_config_json()).expect("parse");
834        assert_eq!(c.num_hidden_layers, 45);
835        assert_eq!(c.layer_types.len(), 45, "text stack stays 0..=44");
836        assert_eq!(c.mtp_layer_types, vec![LayerType::SparseAttention]);
837        // Index 45 resolves, and it resolves through the MTP list, not the text stack.
838        assert_eq!(c.layer_type_at(45), Some(LayerType::SparseAttention));
839        assert_eq!(c.layer_type_at(46), None);
840        assert!(c.has_sparse_attention());
841        assert_eq!(c.sparse_attention_layers().len(), 11, "text stack only");
842    }
843
844    #[test]
845    fn kda_geometry_from_linear_attn_config() {
846        let c = parse_glm5_next(&glm53_config_json()).expect("parse");
847        assert_eq!(c.linear_num_key_heads, 64);
848        assert_eq!(c.linear_key_head_dim, 128);
849        assert_eq!(c.linear_conv_kernel_dim, 4);
850    }
851
852    #[test]
853    fn indexer_topk_is_2048_not_the_deepseek_default() {
854        let c = parse_glm5_next(&glm53_config_json()).expect("parse");
855        assert_eq!(c.index_topk, 2048);
856    }
857
858    #[test]
859    fn missing_rope_key_is_refused_not_guessed() {
860        let mut v: serde_json::Value = serde_json::from_str(&glm53_config_json()).unwrap();
861        v["text_config"]
862            .as_object_mut()
863            .unwrap()
864            .remove("qk_rope_head_dim");
865        let err = parse_glm5_next(&v.to_string()).unwrap_err();
866        assert!(
867            err.to_string().contains("refusing to guess"),
868            "unexpected error: {err}"
869        );
870    }
871
872    /// GLM's SwiGLU clamp is asymmetric and fires only on the activation tails, so a
873    /// checkpoint that declares it and a parser that defaults it disagree SILENTLY on
874    /// well-scaled inputs. The parser must refuse rather than fall back to "no clamp".
875    #[test]
876    fn a_missing_swiglu_limit_is_refused_not_defaulted() {
877        let mut v: serde_json::Value = serde_json::from_str(&glm53_config_json()).unwrap();
878        v["text_config"]
879            .as_object_mut()
880            .unwrap()
881            .remove("swiglu_limit");
882        let err = parse_glm5_next(&v.to_string()).unwrap_err();
883        assert!(
884            err.to_string().contains("swiglu_limit"),
885            "unexpected error: {err}"
886        );
887    }
888
889    #[test]
890    fn the_swiglu_limit_is_read_verbatim() {
891        let c = parse_glm5_next(&glm53_config_json()).unwrap();
892        assert_eq!(c.swiglu_limit, 10.0);
893    }
894
895    /// `n_group`/`topk_group` > 1 would make the group mask load-bearing, and
896    /// `glm5next_router_topk` returns without writing rather than implementing it. Refuse at
897    /// parse time so the failure names the config, not an empty selection row.
898    #[test]
899    fn grouped_expert_routing_is_refused() {
900        for key in ["n_group", "topk_group"] {
901            let mut v: serde_json::Value = serde_json::from_str(&glm53_config_json()).unwrap();
902            v["text_config"][key] = serde_json::json!(8);
903            let err = parse_glm5_next(&v.to_string()).unwrap_err();
904            assert!(
905                err.to_string().contains("Grouped expert routing"),
906                "{key}: unexpected error: {err}"
907            );
908        }
909    }
910
911    #[test]
912    fn contradictory_layer_maps_are_rejected() {
913        let mut v: serde_json::Value = serde_json::from_str(&glm53_config_json()).unwrap();
914        // Claim layer 0 is full attention in layer_types while the index list
915        // says KDA — must not be silently resolved.
916        v["text_config"]["layer_types"][0] =
917            serde_json::Value::String("deepseek_sparse_attention".into());
918        let err = parse_glm5_next(&v.to_string()).unwrap_err();
919        assert!(err.to_string().contains("disagrees"), "unexpected: {err}");
920    }
921
922    // ───────────────────────────────────────────── GLM router dtype ladder (Slice 10)
923
924    /// 🔴 The production default is HF's fp32 router. An absent `moe_router_dtype` means "the
925    /// checkpoint did not say", and the reference implementation's answer for that is fp32 — the
926    /// OPPOSITE of vLLM's fallthrough to bf16.
927    #[test]
928    fn glm_router_defaults_to_hf_fp32_when_the_config_is_silent() {
929        let cfg = parse_glm5_next(&glm53_config_json()).unwrap();
930        assert_eq!(cfg.glm5next_router_mode, Glm5NextRouterMode::HfFp32);
931        assert!(cfg.glm5next_router_mode.is_fp32());
932    }
933
934    /// The compatibility mode must be reachable EXPLICITLY, from its own field — never inferred
935    /// from quantization settings or from the weight-storage precision schedule.
936    #[test]
937    fn glm_router_bf16_compat_mode_is_explicit_and_never_inferred() {
938        let base = glm53_config_json();
939        let with = base.replace(
940            r#""hc_mult": 4,"#,
941            r#""hc_mult": 4, "moe_router_dtype": "bfloat16","#,
942        );
943        assert_ne!(with, base, "fixture anchor moved");
944        let cfg = parse_glm5_next(&with).unwrap();
945        assert_eq!(cfg.glm5next_router_mode, Glm5NextRouterMode::VllmBf16);
946        assert!(!cfg.glm5next_router_mode.is_fp32());
947
948        let fp32 = base.replace(
949            r#""hc_mult": 4,"#,
950            r#""hc_mult": 4, "moe_router_dtype": "float32","#,
951        );
952        assert_eq!(
953            parse_glm5_next(&fp32).unwrap().glm5next_router_mode,
954            Glm5NextRouterMode::HfFp32
955        );
956    }
957
958    /// An unrecognised value must be a hard error. Silently falling back would pick one of two
959    /// semantics at random, which is the whole defect this switch exists to prevent.
960    #[test]
961    fn an_unknown_router_dtype_is_refused_not_defaulted() {
962        let bad = glm53_config_json().replace(
963            r#""hc_mult": 4,"#,
964            r#""hc_mult": 4, "moe_router_dtype": "fp8","#,
965        );
966        assert!(parse_glm5_next(&bad).is_err());
967    }
968}