feat(glm5): add GLM-5.2 model support - #310
Merged
Merged
Conversation
12 tasks
added 4 commits
August 13, 2026 05:23
Add GLM5Config.from_hf()/validate() and GLM52Config(GLM5Config) registered as glm_moe_dsa_5_2 with the GLM-5.2 deltas (1M context, nested rope, DSA indexer reuse fields index_topk_freq/indexer_types). GLM-5 defaults are unchanged. Replace the initializer's hand-parsed model config with the resolved config projected to the minimal engine ModelConfig: head_dim comes from qk_head_dim (not the default), and compressed_kv_dim + first_k_dense_replace are carried. Add a headless regression test locking GLM-5 to its prior resolved values, the head_dim projection, GLM-5.2 reading its real config.json, fail-loud on empty/null/inconsistent config, and the unlisted-variant warning.
Add module-level dsa_layer_skips_topk(config, layer_id) and assert_indexer_schedule_consistent to the GLM-5 config. The helper is value-based: it keys off index_topk_freq presence, so it works on any config object and returns False when the frequency is absent/None/1, leaving GLM-5's uniform per-layer recompute unchanged. The frequency/offset formula is authoritative; the checkpoint indexer_types list is a startup cross-check oracle that fails loud on disagreement.
GLM-5 was the lone model building its graph from an HF transformers.PretrainedConfig; every other model already uses its internal BaseModelConfig dataclass. Migrate GLM-5 to match so GLM-5.2's DSA schedule fields (index_topk_freq/indexer_types) reach the graph without HF plumbing. - config.py: add rope_parameters + rope_type (with from_hf mapping), the only two graph-read attrs the internal config lacked versus the HF class. - glm5_initializer.py: build from the resolved internal config; reuse it in _parse_model_config (no double-resolve); assert the indexer schedule at build (a no-op for GLM-5). - glm5_parameter_server.py: build the graph from the resolved config. - model.py: import the internal GLM5Config for the type hints and the dsa_layer_skips_topk helper for the upcoming indexer-reuse step. - delete configuration_glm5.py.
GLM-5.2 runs its DSA indexer on the full layers only; shared layers carry no indexer weights and reuse the previous full layer's top-k. Prepare the model side of that reuse for the eager (non-CUDA-graph) decode path. - model.py Glm5MLA.__init__: set skip_topk/next_skip_topk; build Glm5Indexer only on full DSA layers, else self.indexer = None (use_dense_mla routing preserved). Glm5Model.__init__ guards the indexer.rotary_emb assignment for None (shared) layers. - wrappers.py: add ClassVar _dsa_prev_topk_indices (the carried top-k) and guard _maybe_init_fused_kernels and the prefill indexer-KV compute on indexer is None. - Parallel_Strategy_Manager.py: guard the indexer FP8 scale attach on None. GLM-5 has no such schedule, so every new branch is dead for it and its behavior is unchanged.
Andrewxu313
force-pushed
the
tairanxu/glm52-model
branch
from
August 12, 2026 21:23
301f639 to
ed8b545
Compare
Andrewxu313
marked this pull request as ready for review
August 12, 2026 21:23
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Add GLM-5.2 model support on top of the config resolver.
GLM5Configgainsfrom_hf()/validate();GLM52Config(GLM5Config)is registered asglm_moe_dsa_5_2with the GLM-5.2 deltas (1M context, nested rope, DSA indexerreuse fields). The initializer projects the resolved config to the engine
ModelConfig(head_dim fromqk_head_dim; carrycompressed_kv_dim+first_k_dense_replace). GLM-5 is migrated off the HFPretrainedConfigtothe internal config end-to-end so the DSA schedule fields reach the graph
without HF plumbing, and a
dsa_layer_skips_topkschedule helper is added. Theeager decode path is made None-safe for GLM-5.2 shared DSA layers
(
skip_topk/next_skip_topk,indexer = Noneon shared layers, a carried_dsa_prev_topk_indices). GLM-5 has no such schedule, so every new branch isdead for it and its behavior is unchanged.
Why
GLM-5.2 shares the GLM-5 model graph but needs its own config identity and a
checkpoint-backed source of truth, plus the model-side hooks for the DSA
indexer top-k reuse used by the shared layers. Stacked on the config seam in
#309 — review/merge that first.
Type of Change
modelkernelcorefixinfradocsFile changes
batchgen/models/glm/glm5/config.pyGLM52Configidentity,from_hf/validate,dsa_layer_skips_topkhelper, rope attrsbatchgen/models/glm/glm5/glm5_initializer.pyModelConfig; assert indexer schedule at buildbatchgen/models/glm/glm5/glm5_parameter_server.pybatchgen/models/glm/glm5/model.pyskip_topk/next_skip_topk; buildGlm5Indexeronly on full DSA layersbatchgen/models/glm/glm5/wrappers.py_dsa_prev_topk_indicescarry; None-guard fused-kernel init + prefill indexer KVbatchgen/models/glm/glm5/Parallel_Strategy_Manager.pyindexer is Nonebatchgen/models/glm/glm5/configuration_glm5.pyPretrainedConfigtests/test_batchgen_model_config.pyChecklist
model/kernelPR does not touch the scheduling/scaffolding layer (§2.5–§2.6).debug_*/scratch_*/tmp_*scripts,BATCHGEN_*env-guards, strayprint(), committed artifacts, orCo-Authored-By(§1/§4). (bench_*benchmarks are fine.)tests/; touched modules'MODULE.mdupdated if the public API changed (§2).bash .github/workflows/scripts/check-pr-hygiene.shlocally (against thetairanxu/glm52-corebase); the CI hygiene check is green.🤖 Generated with Claude Code