Add GLM-5.2 IndexShare DSA, MTP, and GGUF export support - #404
Add GLM-5.2 IndexShare DSA, MTP, and GGUF export support#404justinchuby wants to merge 25 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) |
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. |
| if source_params == (target_bits, target_block_size): | ||
| repacked = repack_gguf_tensor( | ||
| raw_expert.ravel().view(np.uint8), | ||
| qtype_val, | ||
| (n_out, k_in), | ||
| ) | ||
| else: | ||
| _require_supported_requantization( | ||
| bits=target_bits, | ||
| block_size=target_block_size, | ||
| tensor_name=hf_name, | ||
| ) | ||
| values = gguf_model.dequantize_raw_tensor( | ||
| raw_expert, | ||
| qtype, | ||
| (n_out, k_in), | ||
| ) | ||
| repacked = repack_dequantized_tensor( | ||
| values, | ||
| bits=target_bits, | ||
| block_size=target_block_size, | ||
| symmetric=target_symmetric, | ||
| ) | ||
| n_requantized += 1 | ||
|
|
There was a problem hiding this comment.
Fixed in 751645b. The per-expert fast path now checks whether the repacked tensor zero-point presence matches the target symmetric/asymmetric config. A mismatch forces dequantize+requantize, ensuring required zero-point initializers are emitted with the target semantics.
| def _tiny_glm_moe_dsa_config(**overrides): | ||
| defaults = dict( | ||
| model_type="glm_moe_dsa", | ||
| num_hidden_layers=4, | ||
| num_attention_heads=4, |
There was a problem hiding this comment.
Confirmed the branch now includes a representative tiny glm_moe_dsa entry in tests/_test_configs.py, marked representative so config-driven build and alignment suites cover it.
justinchuby
changed the title
590c7da to
fb52e72
Compare
|
♻️ 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>
cd782dd to
87056ad
Compare
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
| qc = config.quantization | ||
| bits = qc.bits | ||
| block_size = qc.group_size |
There was a problem hiding this comment.
Fixed in 751645b by wiring asymmetric quantization through. FusedQuantizedMoE now declares packed FC1/FC2 zero-point parameters when sym=False and passes them to QMoE inputs 11/12; symmetric configs continue using implicit midpoint zero-points.
| The GLM/DeepSeek MoE path decomposes experts into per-expert MLPs, so each | ||
| expert's gate/up/down projections also become MatMulNBits when quantized. | ||
| 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. |
There was a problem hiding this comment.
Updated the docstring to distinguish this decomposed MatMulNBits exporter from the fused QMoE path in export_glm_tiny_qmoe.py.
| OUT = "/home/justinchu/glm-e2e-artifacts/glm-5.2-tiny-q4" | ||
| FP32_OUT = "/home/justinchu/glm-e2e-artifacts/glm-5.2-tiny" |
There was a problem hiding this comment.
Replaced absolute paths with repo-relative artifacts/ defaults, configurable via --output-dir, --fp32-output-dir, or MOBIUS_ARTIFACTS_DIR.
| OUT = "/home/justinchu/glm-e2e-artifacts/glm-5.2-tiny-qmoe" | ||
| FP32_OUT = "/home/justinchu/glm-e2e-artifacts/glm-5.2-tiny" |
There was a problem hiding this comment.
Replaced absolute paths with repo-relative artifacts/ defaults, configurable via --output-dir, --fp32-output-dir, or MOBIUS_ARTIFACTS_DIR.
| from mobius.integrations.onnx_genai import write_inference_metadata | ||
| from mobius.tasks import get_task | ||
|
|
||
| OUT = "/home/justinchu/ds-e2e-artifacts/deepseek-v2-tiny" |
There was a problem hiding this comment.
Replaced the absolute path with a repo-relative artifacts/ default, configurable via --output-dir or MOBIUS_ARTIFACTS_DIR.
| Multi-component packages contain `model/model.onnx`, `mtp/model.onnx`, and an | ||
| `mtp_config.json` sidecar. Both runtime aliases (`onnx-genai` and `ort-genai`) | ||
| emit these artifacts. ORT GenAI does not yet orchestrate the sidecar itself. |
There was a problem hiding this comment.
Chose option (a): onnx-genai now recognizes the exact {model, mtp} package, emits normal decoder metadata plus the same mtp_config.json contract as ORT GenAI, and reports it through the CLI. This was lower risk than weakening the documented cross-runtime artifact contract because it is a narrow exact-shape dispatch with focused tests.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
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"
| qc = config.quantization | ||
| use_fused_qmoe = ( | ||
| config.fused_quantized_moe and qc is not None and qc.quant_method != "none" | ||
| ) | ||
| if use_fused_qmoe: | ||
| self.moe = FusedQuantizedMoE(config, gate=gate) | ||
| else: | ||
| self.moe = MoELayer(config, gate=gate, linear_class=linear_class) |
There was a problem hiding this comment.
Fixed in 751645b. DeepSeek preprocessing now diverts routed-expert tensors into a QMoE packer: fused 3-D GPTQ/AWQ tensors and per-expert GGUF tensors become expert-major FC1/FC2 weights, scales, and zero-points. FC1 gate/up rows are interleaved for swiglu_fusion=1; shared experts remain on the dense path.
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"
| def test_runtime_onnx_genai_calls_write_onnx_genai_config(self): | ||
| """--runtime onnx-genai calls the unified config writer.""" | ||
| with ( |
There was a problem hiding this comment.
Chose option (a): onnx-genai now recognizes the exact {model, mtp} package and emits mtp_config.json; the CLI test verifies the returned sidecar is surfaced. The change is narrowly scoped to that exact component set, so it preserves the existing rejection of unknown multi-component packages.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Forward fused GPTQ g_idx tensors, omit symmetric 8-bit zero points, and pass string paths to ONNX Runtime sessions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
| build_parser.add_argument( | ||
| "--glm-full-attention", | ||
| action="store_true", | ||
| help="Disable GLM-5.2 IndexShare DSA and export the dense MLA fallback.", | ||
| ) |
|
|
Triage/update from Sapper:
The shared PR branch was intentionally not force-updated because the conflict resolution is a large merge and another 404 worktree exists. Merge-ready replacement branch: Justin: please replace/update |
Resolve conflicts with current main while preserving GLM-5.2 IndexShare DSA, MTP, GGUF, and fused QMoE support. Incorporate incoming DeepSeek V4, pipeline metadata, attention decomposition, and CI updates. Address remaining review items for GLM export script docstrings and --glm-full-attention CLI coverage.\n\nCo-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>
|
|
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| @@ -0,0 +1,618 @@ | |||
| # Copyright (c) Microsoft Corporation. | |||
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>
| sym=True, | ||
| ), | ||
| ) | ||
| with pytest.raises(ValueError, match="CUDA.*ignores router_weights"): |
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,
|
@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
Merged Also fixed the two CI linter failures brought in with the last commits (
|
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).
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