feat(glm5): add GLM-5.2 model support

batchgen/models/glm/glm5/Parallel_Strategy_Manager.py

@@ -688,8 +688,9 @@ def _setup_fp8_scales(self):
             # After wrapping, self_attn is GLM5AttnWrapper; original Glm5MLA is at .module
             inner = attn.module if hasattr(attn, 'module') else attn
             # When use_dense_mla is set, Glm5MLA skips indexer construction;
-            # skip the scale attach too (no destination).
-            if hasattr(inner, "indexer"):
+            # skip the scale attach too (no destination). GLM-5.2 "shared" layers
+            # carry no indexer weights either (indexer is None) — skip them.
+            if getattr(inner, "indexer", None) is not None:
                 indexer = inner.indexer
                 for proj, attr in [("wk", "wk_scale"), ("wq_b", "wq_b_scale")]:
                     key = f"model.layers.{layer_idx}.self_attn.indexer.{proj}.weight_scale_inv"

batchgen/models/glm/glm5/glm5_initializer.py

@@ -22,7 +22,8 @@
 import torch
 
 from batchgen.config.config import EngineConfig, ModelConfig
-from .configuration_glm5 import Glm5Config
+from batchgen.config.batchgen_model_config import BatchGenModelConfig
+from .config import assert_indexer_schedule_consistent
 from .set_basic_config import set_basic_config
 from .planner import GLM5Planner
 from batchgen.kv_cache.host_kv_mananger_config import build_host_kv_config
@@ -36,9 +37,23 @@
 
 class GLM5Initializer:
     def __init__(self, input_arguments):
-        self.loaded_model_config = Glm5Config()
-        self.loaded_model_config._name_or_path = input_arguments.huggingface_ckpt_name
+        self.model_name = input_arguments.huggingface_ckpt_name
+        # Local checkpoint dir holding config.json (pre-downloaded model files).
+        # None for a bare HF id -> resolver falls back to rich defaults.
+        self.checkpoint_path = input_arguments.get("cache_dir", None)
+
+        # Single source of truth: the resolved internal config (GLM5Config /
+        # GLM52Config, checkpoint-backed). This IS the config used to build the
+        # model graph (loaded_model_config) — GLM-5 no longer uses an HF
+        # transformers.PretrainedConfig, matching kimi and every other model.
+        self.loaded_model_config = BatchGenModelConfig.resolve(
+            self.model_name, self.checkpoint_path
+        )
+        self.loaded_model_config._name_or_path = self.model_name
         self.loaded_model_config.architectures = ["GlmMoeDsaForCausalLM"]
+        # Fail loud at build if the DSA indexer schedule (freq/offset) disagrees
+        # with the checkpoint's per-layer indexer_types (no-op for GLM-5).
+        assert_indexer_schedule_consistent(self.loaded_model_config)
 
         self.host_kv_cache_size = input_arguments.host_kv_cache_size
         self.host_kv_cache_byte_size = input_arguments.host_kv_cache_size * (1024**3)
@@ -156,15 +171,35 @@ def _default_engine_config(self):
         }
 
     def _parse_model_config(self):
+        """Resolve the rich, checkpoint-backed config via the single decoupled
+        resolver, then PROJECT it into the minimal engine ModelConfig.
+
+        The rich config (GLM5Config / GLM52Config) is the single source of
+        truth; the engine's minimal ModelConfig is duck-typed on the C++ side
+        (core/utils.cpp parse_model_config reads model_type, num_hidden_layers,
+        num_local_experts, num_attention_heads, num_key_value_heads, head_dim).
+
+        Traps honored here:
+        * head_dim <- rich.qk_head_dim (256), NOT rich.head_dim (64). The C++
+          attention geometry expects the MLA qk head dim.
+        * compressed_kv_dim and first_k_dense_replace are NOT declared fields of
+          the minimal ModelConfig but ARE read downstream (this file @
+          _default_engine_config, glm5_parameter_server). Set them as attributes
+          so Init does not crash with AttributeError.
+        """
+        rich = self.loaded_model_config
+
         model_config = ModelConfig()
-        model_config.model_type = "glm_moe_dsa"
-        model_config.num_hidden_layers = 78
-        model_config.num_local_experts = 256
-        model_config.num_attention_heads = 64
-        model_config.num_key_value_heads = 64
-        model_config.head_dim = 256  # qk_nope + qk_rope = 192 + 64
-        model_config.compressed_kv_dim = 576  # kv_lora_rank + qk_rope_head_dim
-        model_config.first_k_dense_replace = 3
+        model_config.model_type = rich.model_type
+        model_config.num_hidden_layers = rich.num_hidden_layers
+        model_config.num_local_experts = rich.num_local_experts
+        model_config.num_attention_heads = rich.num_attention_heads
+        model_config.num_key_value_heads = rich.num_key_value_heads
+        # TRAP: engine head_dim is the MLA qk head dim (256), not rich.head_dim (64).
+        model_config.head_dim = rich.qk_head_dim
+        # Not declared on ModelConfig but read downstream — set as attributes.
+        model_config.compressed_kv_dim = rich.compressed_kv_dim  # kv_lora_rank + qk_rope_head_dim
+        model_config.first_k_dense_replace = rich.first_k_dense_replace
         return model_config
 
     def Init(self, weights_storage):

batchgen/models/glm/glm5/glm5_parameter_server.py

@@ -38,11 +38,9 @@ def _diag(msg):
 from tqdm import tqdm, trange
 _diag("tqdm done")
 
-from .configuration_glm5 import Glm5Config
-_diag("configuration_glm5 done")
 from .model import Glm5ForCausalLM
 _diag("model (Glm5ForCausalLM) done")
-from batchgen.config.model_registry import load_config
+from batchgen.config.batchgen_model_config import BatchGenModelConfig
 _diag("model_registry done")
 
 try:
@@ -64,8 +62,11 @@ def __init__(self, huggingface_ckpt_name, cache_dir, converted_ckpt_dir, enable_
         self.state_dict_name_map = {}
         self.enable_hugetlbfs = enable_hugetlbfs
         self.enable_memfd = enable_memfd
-        self.model_config = load_config(huggingface_ckpt_name)
-        self.hf_config = Glm5Config()
+        # Single resolved internal config (checkpoint-backed GLM5Config/GLM52Config).
+        # Used both for metadata reads and to build the model graph — GLM-5 no
+        # longer uses an HF transformers.PretrainedConfig.
+        self.model_config = BatchGenModelConfig.resolve(huggingface_ckpt_name, cache_dir)
+        self.hf_config = self.model_config
         self.hf_config._name_or_path = huggingface_ckpt_name
 
         free_memory, total_memory = torch.cuda.mem_get_info()

batchgen/models/glm/glm5/model.py

@@ -29,7 +29,11 @@
 import torch.nn.functional as F
 
 from .decode_utils import clamp_token_indices_to_seqlens
-from .configuration_glm5 import Glm5Config
+# GLM-5 uses BatchGen's internal config (a plain BaseModelConfig dataclass), not
+# an HF transformers.PretrainedConfig — matching every other model (kimi, etc.).
+# Imported under the historical name `Glm5Config` so the __init__ type hints below
+# need no churn; only attribute reads are used, so any config object works.
+from .config import GLM5Config as Glm5Config, dsa_layer_skips_topk
 
 
 # ============================================================================
@@ -754,11 +758,17 @@ def __init__(self, config: Glm5Config, layer_idx: int):
         # Softmax scale
         self.softmax_scale = self.q_head_dim ** -0.5
 
-        # DSA indexer — structurally absent when config.use_dense_mla is True.
-        # Downstream decode dispatch uses hasattr(self.module, 'indexer') as
-        # the signal to route to the dense-MLA (DeepSeek-V3 / Kimi style) path.
+        # DSA indexer — structurally absent when config.use_dense_mla is True
+        # (the `indexer` attribute is then NOT set, so hasattr(...) routes to the
+        # dense-MLA path — unchanged for DeepSeek-V3 / Kimi style configs).
+        # For DSA configs the attribute always exists, but GLM-5.2 "shared" layers
+        # carry no indexer weights: they reuse the previous full layer's top-k
+        # indices, so self.indexer is None on those layers (GLM-5: never a shared
+        # layer, so indexer is built on every layer — bit-identical).
+        self.skip_topk = dsa_layer_skips_topk(config, layer_idx)
+        self.next_skip_topk = dsa_layer_skips_topk(config, layer_idx + 1)
         if not getattr(config, "use_dense_mla", False):
-            self.indexer = Glm5Indexer(config, layer_idx)
+            self.indexer = None if self.skip_topk else Glm5Indexer(config, layer_idx)
 
         # Absorbed projections for decode (set by initialize())
         self.q_absorb = None
@@ -2298,7 +2308,9 @@ def __init__(self, config: Glm5Config):
         # Assign shared RoPE to attention and indexer
         for layer in self.layers:
             layer.self_attn.rotary_emb = self._shared_rotary_emb
-            if hasattr(layer.self_attn, 'indexer'):
+            # DSA layers have an `indexer` attribute; GLM-5.2 shared layers set it
+            # to None (no indexer weights), so guard the deref.
+            if getattr(layer.self_attn, 'indexer', None) is not None:
                 layer.self_attn.indexer.rotary_emb = self._shared_rotary_emb
 
         self.norm = Glm5RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

batchgen/models/glm/glm5/wrappers.py

@@ -413,6 +413,10 @@ class GLM5AttnWrapper(AttnWrapperBase):
     # Set once per decode step by the worker so per-layer _forward_decode_dsa
     # branches on it without doing a D2H .sum().item() 78 times per step.
     _dsa_short_count: ClassVar[Optional[int]] = None
+    # DSA top-k index REUSE (GLM-5.2): carried top-k indices from the most recent
+    # FULL layer, reused by subsequent shared layers. Reset per decode step by the
+    # worker. For GLM-5 (all layers full) this is never read by a shared branch.
+    _dsa_prev_topk_indices: ClassVar[Optional[torch.Tensor]] = None
     # Whole-model CUDA graph can pad local rows to a global NCCL bucket. These
     # graph-owned overrides let GLM-5 DSA use explicit slot sentinels for padded
     # rows instead of deriving slot count from cur_batch.
@@ -547,7 +551,7 @@ def initialize_fused_kernels(self):
         AND after _setup_fp8_scales has attached indexer.wk_scale / wq_b_scale.
         """
         attn = self.module
-        if not hasattr(attn, "indexer"):
+        if not hasattr(attn, "indexer") or attn.indexer is None:
             return
         indexer = attn.indexer
 
@@ -700,15 +704,20 @@ def _forward_prefill(self, hidden_states: torch.Tensor, **kwargs) -> Tuple:
                 raise RuntimeError(
                     "GLM-5 DSA prefill requires indexer KV; refusing primary-only host offload"
                 )
-            indexer_kv = self.module.indexer.compute_indexer_kv(
-                hidden_states_2d.unsqueeze(0),
-                positions=self.position_ids.to(hidden_states_2d.device),
-            )
-            if indexer_kv is None:
-                raise RuntimeError(
-                    "GLM-5 DSA prefill indexer returned no KV; refusing primary-only host offload"
+            # Shared layers (GLM-5.2) have indexer is None: they reuse a full
+            # layer's top-k at decode and never read an aux indexer cache, so
+            # skip the indexer-K compute + offload entirely. GLM-5 layers always
+            # have a real indexer so this block always runs there.
+            if self.module.indexer is not None:
+                indexer_kv = self.module.indexer.compute_indexer_kv(
+                    hidden_states_2d.unsqueeze(0),
+                    positions=self.position_ids.to(hidden_states_2d.device),
                 )
-            self._offload_prepacked_indexer_kv(indexer_kv.squeeze(0))
+                if indexer_kv is None:
+                    raise RuntimeError(
+                        "GLM-5 DSA prefill indexer returned no KV; refusing primary-only host offload"
+                    )
+                self._offload_prepacked_indexer_kv(indexer_kv.squeeze(0))
 
             self._offload_prepacked_kv(offload_kv)
             attn_output = attn_output.unsqueeze(0)