Add GLM-5.2 IndexShare DSA, MTP, and GGUF export support - #404
Draft
justinchuby wants to merge 28 commits into
Draft
Add GLM-5.2 IndexShare DSA, MTP, and GGUF export support#404justinchuby wants to merge 28 commits into
justinchuby wants to merge 28 commits into
Conversation
Performance Comparison
|
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Contributor
There was a problem hiding this comment.
Pull request overview
Adds first-pass export support for GLM-5.2 / GLM-MoE DSA by wiring the architecture into the registry, enabling GGUF (glm-dsa) config/weight ingestion (including split MLA KV-B fusion and expert-by-expert streaming), and extending the quantized import path to cover additional supported layouts. Also accepts --runtime onnx-genai as an alias for ort-genai.
Changes:
- Register
glm_moe_dsa(HF) /glm-dsa(GGUF) and add a GLM-specific model wrapper reusing the DeepSeek-V3 MLA+MoE backbone with dense-attention fallback. - Extend GGUF import: map GLM tensor names, extract GLM metadata into
ArchitectureConfig, fuse split MLA KV-B weights, and stream stacked expert tensors. - Extend quantized repacking/CLI: support float requantization to 8-bit block-32 targets and accept
--runtime onnx-genai.
Reviewed changes
Copilot reviewed 19 out of 19 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/cli_test.py | Adds CLI coverage for --runtime onnx-genai alias. |
| src/mobius/main.py | Implements onnx-genai as an alias for ort-genai. |
| src/mobius/models/glm_moe_dsa.py | Adds GLM-5.2 DSA model wrapper over DeepSeek MLA+MoE with weight filtering. |
| src/mobius/models/glm_moe_dsa_test.py | Adds GLM-specific graph/config/quantization/preprocess tests. |
| src/mobius/models/deepseek.py | Adds quantized linear/embedding injection via linear_class + QuantizedEmbedding option. |
| src/mobius/components/_deepseek_mla.py | Threads linear_class through MLA projections. |
| src/mobius/components/_moe.py | Threads linear_class into expert MLPs. |
| src/mobius/models/init.py | Exports GlmMoeDsaCausalLMModel. |
| src/mobius/_registry.py | Registers glm_moe_dsa and adds a default model id. |
| src/mobius/_configs/_base.py | Adds DSA/MTP-related config fields and enables rope_interleave default for glm_moe_dsa. |
| src/mobius/integrations/gguf/_tensor_mapping.py | Adds GLM-DDA tensor mapping + skips indexer/nextn tensors. |
| src/mobius/integrations/gguf/_tensor_mapping_test.py | Tests GLM tensor mapping behavior. |
| src/mobius/integrations/gguf/_config_mapping.py | Adds GLM metadata key mapping + postprocess to derive MLA/DSA fields. |
| src/mobius/integrations/gguf/_reader_test.py | Tests GGUF arch→model_type mapping and GLM config extraction. |
| src/mobius/integrations/gguf/_builder.py | Adds split KV-B fusion, expert-by-expert repack, and allows 8-bit requant targets. |
| src/mobius/integrations/gguf/_builder_test.py | Tests GLM normalization behaviors and Q8 requantized head behavior. |
| src/mobius/integrations/gguf/_repacker.py | Extends float requantization to bits=8, block=32. |
| src/mobius/integrations/gguf/_repacker_test.py | Adds Q8 float requantization round-trip bound test. |
| docs/research/glm52-export.md | Documents GLM-5.2 verified architecture, current scope, and follow-ups. |
justinchuby
changed the title
justinchuby
force-pushed
the
glm5.2-moe-export
branch
from
July 16, 2026 15:35
590c7da to
fb52e72
Compare
Member
Author
|
♻️ Rebased onto |
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The IndexShare DSA indexer shares the model's rotary_emb cos/sin cache, which is sized qk_rope_head_dim/2. GlmMoeDsaIndexer.project_key passed rotary_embedding_dim=index_head_dim//2 to the opset-24 RotaryEmbedding op, which then expects a cos cache of (index_head_dim//2)/2 and fails shape validation at runtime: Input 'cos_cache' dimension 2 should be same as head_size / 2 or rotary_embedding_dim / 2, got 4 The indexer key spans the full index_head_dim and must be fully rotated, matching the main MLA q_rope/k_rope calls (rotary_embedding_dim=0). Fixes end-to-end inference of glm_moe_dsa models through the onnx-genai engine. Adds a regression test asserting indexer RotaryEmbedding nodes use full rotation (rotary_embedding_dim=0). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Sibling of export_glm_tiny.py. Attaches an int4 block-32 asymmetric QuantizationConfig so the tiny glm_moe_dsa linear projections and per-expert MoE MLPs are emitted as com.microsoft::MatMulNBits, producing an onnx-genai-loadable artifact for the quantized E2E smoke test. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…experts Add FusedQuantizedMoE component that emits a single int4/block-32 expert-major com.microsoft::QMoE node (swiglu_fusion=1) for the shared GLM/DeepSeek routed expert stack, gated behind ArchitectureConfig.fused_quantized_moe. GLM sigmoid + noaux_tc routing is preserved exactly: scores_for_choice feeds QMoE router_probs for top-k selection, and the normalized+scaled combine weights are scattered into the optional router_weights input (normalize_routing_weights=0). Adds DeepSeekMoEGate.route_for_qmoe, export_glm_tiny_qmoe.py, and a fused-path unit test. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
justinchuby
force-pushed
the
glm5.2-moe-export
branch
from
July 21, 2026 22:47
cd782dd to
87056ad
Compare
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 30 out of 30 changed files in this pull request and generated 1 comment.
Comments suppressed due to low confidence (5)
src/mobius/components/_moe.py:296
- FusedQuantizedMoE is documented as symmetric-only (no zero-points), but init does not enforce QuantizationConfig.sym=True. If a caller enables fused_quantized_moe with asymmetric quantization, the graph will still be built but will pass None for *zero_points inputs to QMoE, which is inconsistent with the quantization setup.
qc = config.quantization
bits = qc.bits
block_size = qc.group_size
if bits not in (1, 2, 4, 8):
raise ValueError(f"QMoE expert_weight_bits must be 1/2/4/8, got {bits}")
if block_size < 16 or (block_size & (block_size - 1)):
raise ValueError(f"QMoE block_size must be a power of 2 >= 16, got {block_size}")
export_glm_tiny_quant.py:11
- This docstring states mobius has no fused com.microsoft::QMoE emitter and that QMoE is not emitted, but this PR introduces FusedQuantizedMoE and export_glm_tiny_qmoe.py. The docstring is now inaccurate and will mislead readers.
mobius has no fused ``com.microsoft::QMoE`` emitter for any model (only an
unquantized ``com.microsoft::MoE`` path in gemma4), so QMoE is NOT emitted
here — see the decision note for the scoping gap.
export_glm_tiny_quant.py:36
- Hard-coded absolute output paths make this script non-portable (and leak a developer-specific directory layout into the repo). Prefer environment variables and a safe default like /tmp so contributors can run it without editing source.
OUT = "/home/justinchu/glm-e2e-artifacts/glm-5.2-tiny-q4"
FP32_OUT = "/home/justinchu/glm-e2e-artifacts/glm-5.2-tiny"
export_glm_tiny_qmoe.py:36
- Hard-coded absolute output paths make this script non-portable. Prefer environment variables and a safe default like /tmp so contributors can run it without editing source.
OUT = "/home/justinchu/glm-e2e-artifacts/glm-5.2-tiny-qmoe"
FP32_OUT = "/home/justinchu/glm-e2e-artifacts/glm-5.2-tiny"
export_deepseek_v2_tiny.py:19
- Hard-coded absolute output path makes this script non-portable. Prefer an environment variable and a safe default like /tmp so contributors can run it without editing source.
OUT = "/home/justinchu/ds-e2e-artifacts/deepseek-v2-tiny"
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 30 out of 30 changed files in this pull request and generated 2 comments.
Comments suppressed due to low confidence (4)
src/mobius/models/deepseek.py:328
fused_quantized_moeswitches the MoE implementation toFusedQuantizedMoE, which expects packed expert-major tensors (fc1_experts_weights,fc1_scales,fc2_experts_weights,fc2_scales) instead of the existing per-expert projections. The current GGUF loader streams routed expert tensors into per-expert keys (e.g....mlp.experts.{i}.gate_proj.weight), and existing preprocessors map those under...mlp.moe.experts.*, so enabling this flag is likely to make weight application fail due to missing packed initializers. Either add a packing step (similar to Gemma4’s fused-expert preprocessing) or raise a clear error when packed tensors aren’t provided.
qc = config.quantization
use_fused_qmoe = (
config.fused_quantized_moe and qc is not None and qc.quant_method != "none"
)
if use_fused_qmoe:
export_glm_tiny_quant.py:11
- The module docstring claims mobius has no fused
com.microsoft::QMoEemitter, but this PR addsFusedQuantizedMoEand afused_quantized_moeconfig flag. This comment is now misleading for readers using these export helpers.
mobius has no fused ``com.microsoft::QMoE`` emitter for any model (only an
unquantized ``com.microsoft::MoE`` path in gemma4), so QMoE is NOT emitted
here — see the decision note for the scoping gap.
export_glm_tiny_quant.py:36
- Hard-coded absolute output paths (
/home/justinchu/...) make this script non-portable and will fail in most environments/CI. Consider deriving these from CLI args or environment variables with a safe default (e.g., relative to the current working directory).
OUT = "/home/justinchu/glm-e2e-artifacts/glm-5.2-tiny-q4"
FP32_OUT = "/home/justinchu/glm-e2e-artifacts/glm-5.2-tiny"
export_glm_tiny_qmoe.py:36
- Hard-coded absolute output paths (
/home/justinchu/...) make this script non-portable and will fail in most environments/CI. Consider deriving these from CLI args or environment variables with a safe default (e.g., relative to the current working directory).
OUT = "/home/justinchu/glm-e2e-artifacts/glm-5.2-tiny-qmoe"
FP32_OUT = "/home/justinchu/glm-e2e-artifacts/glm-5.2-tiny"
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
|
|
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 43 out of 43 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (4)
src/mobius/models/glm_moe_dsa.py:294
- The PR description states the DSA path is exported using only standard ONNX ops (no custom op required for correctness), but this implementation emits a custom
pkg.nxrt::IndexSharenode. That makes the exported model dependent on a non-standard opset and contradicts the stated portability/correctness claim.
Please either (a) add a portable path that lowers DSA to standard ops (e.g., dense additive mask + Attention) behind a config/CLI flag, or (b) update the PR description/docs to explicitly call out the custom-op dependency for DSA builds.
docs/research/glm52-export.md:88
- This section claims the exported DSA graph uses only standard ONNX ops and requires no custom op for correctness, but the implementation and tests in this PR emit a custom
pkg.nxrt::IndexShareop (and assertAttentioncount is 0). The doc should be updated to match the current export contract.
Portable ONNX preserves DSA selection numerics by converting the selected
indices to an additive mask for the standard `Attention` operator. No custom op
is required for correctness. Practical million-token performance still needs
a selected-token sparse-attention kernel so the runtime avoids evaluating
masked K/V positions.
src/mobius/rewrite_rules/_group_query_attention_test.py:206
_build_tiny_mla_graph()is introduced here but is not referenced anywhere else in this test module. Keeping unused helpers in test code makes future refactors harder and can confuse readers about intended coverage.
Either wire this helper into a test case, or remove it if it’s leftover from earlier iterations.
src/mobius/rewrite_rules/_decompose_attention_test.py:172
- The acceptable error tolerance for the softcap cases was loosened from
1e-5to5e-3, which is a large relaxation and risks masking real parity regressions in the attention decomposition.
Please add a short comment explaining why this larger atol is required (e.g., known numerical instability/EP differences), and consider tightening it as much as possible.
Keep fused QMoE activation inputs in the model dtype and fail fast for CUDA exports that require router_weights, since ORT CUDA QMoE currently ignores input 14. Add regression tests for activation dtype and the CUDA guard.\n\nCo-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 43 out of 43 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (2)
src/mobius/main.py:291
- In the
--modelcode path,--glm-full-attentionalways passesconfig_overridesintobuild(), so it will also disable DSA/MTP for any model type, not justglm_moe_dsa. This is surprising for a GLM-specific flag and can silently alter other architectures that use these fields. Consider loading the HF config when the flag is present, verifyingmodel_type == 'glm_moe_dsa', and only then applying overrides.
config_overrides=(
{"use_dsa": False, "num_nextn_predict_layers": 0}
if args.glm_full_attention
else None
),
src/mobius/main.py:251
--glm-full-attentionis applied unconditionally when building from--config, even if the resolvedmodel_typeis notglm_moe_dsa. Sinceuse_dsa/num_nextn_predict_layersalso affect other architectures (e.g. DeepSeek-V4 CSA/MTP), this flag can silently change unrelated exports. Please validate the model type and fail fast when the flag is used on non-GLM-5.2 configs.
This issue also appears on line 287 of the same file.
if args.glm_full_attention:
config = dataclasses.replace(
config,
use_dsa=False,
num_nextn_predict_layers=0,
justinchuby
marked this pull request as ready for review
July 30, 2026 04:57
Member
Author
|
@copilot merge from main |
Signed-off-by: GitHub <noreply@github.com> # Conflicts: # src/mobius/integrations/gguf/_tensor_mapping.py # src/mobius/integrations/onnx_genai/auto_export.py
Contributor
Merged Also fixed the two CI linter failures brought in with the last commits (
|
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 43 out of 43 changed files in this pull request and generated no new comments.
Comments suppressed due to low confidence (4)
src/mobius/models/glm_moe_dsa.py:287
- The PR description claims Portable ONNX correctness for DSA uses only standard ops and that no custom op is required, but this implementation hard-depends on the custom
pkg.nxrt::IndexShareop for the DSA path. That means a generic ONNX Runtime / non-nxrt environment cannot run the exported graph as-is.
Either (a) add a portable standard-op fallback for DSA (e.g., dense Attention with a ScatterElements-built sparse mask) behind a config flag, or (b) update the PR description/docs to state that DSA execution requires the pkg.nxrt custom-op runtime.
src/mobius/models/glm_moe_dsa.py:576
GlmMoeDsaCausalLMModel.__init__only computes/validates the IndexShare schedule whenconfig.indexer_types is None. If a caller provides an explicitindexer_typeslist with the wrong length, the model will accept it and later code (task cache packing, per-layer indexing) can fail with anIndexErrorinstead of the intended clearValueError.
Consider validating whenever DSA is enabled (e.g., if config.use_dsa: config.indexer_types = _indexer_types(config)), since _indexer_types() already performs the length check for caller-supplied schedules.
src/mobius/models/glm_moe_dsa.py:26
- This file imports
apply_rotary_pos_embfrom a private module path (mobius.components._rotary_embedding). Project convention is to import components/helpers from the publicmobius.componentsAPI to avoid brittle private-module dependencies.
Consider re-exporting apply_rotary_pos_emb from mobius.components and importing it from there (or provide a small local wrapper) so the new model doesn’t depend on a private module path.
src/mobius/rewrite_rules/_decompose_attention_test.py:172
- The parity tolerance for the softcap cases was loosened from
1e-5to5e-3, which is a large change and can mask regressions in the softcap path. If this is required due to a known ORT/kernel numeric change, please add a short comment explaining why this tolerance is expected/acceptable (and ideally link to an issue or ORT version note).
# Conflicts: # src/mobius/_builder.py # src/mobius/integrations/ort_genai/auto_export.py
…metric-MoE zero-points, docstrings, --glm-full-attention help Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 43 out of 43 changed files in this pull request and generated no new comments.
Suppressed comments (3)
src/mobius/integrations/onnx_genai/auto_export.py:577
- In the special-case
['model', 'mtp']branch,write_decoder_metadata()is called but_add_explicit_io_to_file()is not. Every other pipeline branch (and the single-decoder fallback) augments the emitted YAML with explicit I/O derived from the actual ONNX ports; skipping it here can leaveinference_metadata.yamlmissing/incorrectinputs/outputsrole data for the package.
if component_names == ["model", "mtp"]:
path = write_decoder_metadata(
output_dir, config=resolved_config, kv_native_dtype=kv_native_dtype
)
artifacts = {
"inference_metadata": path,
"mtp_config": _write_mtp_config(output_dir, resolved_config),
}
src/mobius/models/glm_moe_dsa.py:569
- The class docstring says "portable IndexShare", but the DSA path relies on a custom-domain op (
pkg.nxrt::IndexShare). "Portable" is misleading here unless there is also a standard-op fallback.
docs/research/glm52-export.md:88 - This section claims the DSA path is "Portable ONNX" with only standard ops and "No custom op". However, the implementation currently emits the custom
pkg.nxrt::IndexShareoperator (and explicitly avoidsAttention/ScatterElements) for the DSA attention path, so a runtime implementation of that custom op is required.
Portable ONNX preserves DSA selection numerics by converting the selected
indices to an additive mask for the standard `Attention` operator. No custom op
is required for correctness. Practical million-token performance still needs
a selected-token sparse-attention kernel so the runtime avoids evaluating
masked K/V positions.
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 43 out of 43 changed files in this pull request and generated no new comments.
Suppressed comments (2)
docs/research/glm52-export.md:88
- This section says the DSA path is implemented by converting selected indices into an additive mask for the standard
Attentionop and that no custom op is required. The current implementation lowers sparse attention to the custompkg.nxrt::IndexShareop (and only the dense fallback uses standardAttention), so this paragraph is misleading and contradicts the exported graphs.
Portable ONNX preserves DSA selection numerics by converting the selected
indices to an additive mask for the standard `Attention` operator. No custom op
is required for correctness. Practical million-token performance still needs
a selected-token sparse-attention kernel so the runtime avoids evaluating
masked K/V positions.
src/mobius/integrations/onnx_genai/auto_export.py:577
- In the
{model, mtp}package case you returnwrite_decoder_metadata(...)without calling_add_explicit_io_to_file(...), unlike the other code paths. That means the emittedinference_metadata.yamlwon’t be augmented with the actual ONNX input/output port roles even when real models are present, which is inconsistent with the rest of the emitter and can lead to stale metadata if the decoder I/O contract changes.
if component_names == ["model", "mtp"]:
path = write_decoder_metadata(
output_dir, config=resolved_config, kv_native_dtype=kv_native_dtype
)
artifacts = {
"inference_metadata": path,
"mtp_config": _write_mtp_config(output_dir, resolved_config),
}
tokenizer_path = _write_hf_tokenizer(output_dir, source)
if tokenizer_path is not None:
artifacts["tokenizer"] = tokenizer_path
return artifacts
justinchuby
marked this pull request as draft
August 4, 2026 17:00
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
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.
Summary
glm_moe_dsaand llama.cpp/GGUFglm-dsamtp/model.onnx, including embedding/hidden fusion, DSA+MoE block, and shared-head normmtp_config.jsonfor--runtime onnx-genai/ort-genai--glm-full-attentionRuntime status
Portable ONNX preserves DSA selection numerics with standard ops (
MatMul,LayerNormalization,Relu,TopK,ScatterElements, andAttention). No custom op is required for correctness. A selected-token sparse-attention kernel is still needed to realize GLM-5.2's long-context FLOP reduction instead of evaluating masked positions.The MTP graph and step-0
topk_indicesare exported. ORT GenAI does not currently expose separate MLA/indexer cache lengths or native MTP sidecar orchestration, soindex_share_for_mtp_iterationremains a runtime follow-up; later speculative iterations recompute selection rather than silently reusing incorrect cache state.IQ1/IQ2/IQ3 preservation remains a separate workstream and is unchanged here.
Validation
160 passed279 passed9 passed, 2 skipped4509 passed, 329 skipped, 60 xfailedonnx.checker, load in ONNX Runtime CPU, and execute one-token prefill with empty cachesgit diff --checkpassFollow-ups