Add GLM-5.2 IndexShare DSA, MTP, and GGUF export support

.github/workflows/gpu_l4_golden_parity.yml

@@ -59,7 +59,8 @@ jobs:
           pip install -r requirements/ci/requirements.txt
           pip install soundfile librosa ml_dtypes flatbuffers numpy packaging protobuf sympy coloredlogs
           pip install -e '.[testing,transformers]'
-          pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ onnxruntime-gpu onnxruntime-genai-cuda
+          # ORT 1.27+ PyPI wheels require CUDA 13; these runners and PyTorch use CUDA 12.8.
+          pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ 'onnxruntime-gpu<1.27' onnxruntime-genai-cuda
 
       - name: Check for golden test data
         id: check_golden

.github/workflows/gpu_l5_generation_e2e.yml

@@ -59,7 +59,8 @@ jobs:
           pip install -r requirements/ci/requirements.txt
           pip install soundfile librosa ml_dtypes flatbuffers numpy packaging protobuf sympy coloredlogs
           pip install -e '.[testing,transformers]'
-          pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ onnxruntime-gpu onnxruntime-genai-cuda
+          # ORT 1.27+ PyPI wheels require CUDA 13; these runners and PyTorch use CUDA 12.8.
+          pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ 'onnxruntime-gpu<1.27' onnxruntime-genai-cuda
 
       - name: Run L5 generation E2E tests
         env:

.github/workflows/main.yml

@@ -283,6 +283,7 @@ jobs:
           pip install -r requirements/ci/requirements.txt
           pip install soundfile librosa ml_dtypes flatbuffers numpy packaging protobuf sympy coloredlogs
           pip install -e '.[testing]'
+          # ORT 1.27+ PyPI wheels require CUDA 13; these runners and PyTorch use CUDA 12.8.
           pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ "onnxruntime-gpu<1.27" onnxruntime-genai-cuda
       - name: Run fast integration tests
         env:

.github/workflows/validation_examples_gpu.yml

@@ -83,7 +83,8 @@ jobs:
           pip install -r requirements/ci/requirements.txt
           pip install -e '.[testing,transformers]'
           pip install soundfile librosa ml_dtypes flatbuffers numpy packaging protobuf sympy coloredlogs
-          pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ onnxruntime-gpu onnxruntime-genai-cuda
+          # ORT 1.27+ PyPI wheels require CUDA 13; these runners and PyTorch use CUDA 12.8.
+          pip install --pre --extra-index-url https://aiinfra.pkgs.visualstudio.com/PublicPackages/_packaging/ort-cuda-12-nightly/pypi/simple/ 'onnxruntime-gpu<1.27' onnxruntime-genai-cuda
 
       - name: Run ${{ matrix.name }}
         env:

docs/research/glm52-export.md

@@ -0,0 +1,147 @@
+# GLM-5.2 export design
+
+## Verified architecture
+
+Sources:
+
+- [`zai-org/GLM-5.2` config](https://huggingface.co/zai-org/GLM-5.2/blob/main/config.json)
+- Transformers `GlmMoeDsaForCausalLM` implementation
+- llama.cpp `glm-dsa` GGUF metadata and tensor names
+- [`unsloth/GLM-5.2-GGUF`](https://huggingface.co/unsloth/GLM-5.2-GGUF)
+- [IndexShare paper, arXiv:2603.12201](https://arxiv.org/abs/2603.12201)
+- [GLM-5 report, arXiv:2602.15763](https://arxiv.org/abs/2602.15763)
+
+The Hugging Face model type is `glm_moe_dsa`, with
+`GlmMoeDsaForCausalLM` as its architecture. llama.cpp writes the exact
+`general.architecture` value `glm-dsa`.
+
+| Field | GLM-5.2 value |
+| --- | --- |
+| Backbone layers | 78 |
+| MTP layers | 1 additional NextN layer |
+| Hidden / dense FFN size | 6144 / 12288 |
+| Attention heads / KV heads | 64 / 64 |
+| QK / NoPE / RoPE / V head dimensions | 256 / 192 / 64 / 256 |
+| Q-LoRA / KV-LoRA rank | 2048 / 512 |
+| Routed experts / experts per token | 256 / 8 |
+| Routed / shared-expert FFN size | 2048 / 2048 |
+| Shared experts | 1 |
+| Dense prefix | first 3 layers |
+| Router | float32 sigmoid, `noaux_tc`, normalized top-k, scale 2.5 |
+| Norm / activation | pre-RMSNorm, epsilon 1e-5 / SiLU |
+| RoPE | interleaved, theta 8,000,000 |
+| Maximum context | 1,048,576 |
+| Indexer | 32 heads, head dimension 128, top-k 2048 |
+
+GGUF reports `block_count=79` and `nextn_predict_layers=1`; therefore the
+causal-LM backbone has 78 layers. GGUF also splits the MLA KV decompression
+weight into `attn_k_b` and `attn_v_b`, while Transformers exposes one
+`kv_b_proj`.
+
+## Implemented scope
+
+The exporter now preserves the 78-layer backbone, IndexShare DSA, and the
+additional improved-MTP layer. Full indexers are emitted for layers 0, 1, 2,
+6, 10, ... 74; each full layer's INT64 top-k result is reused by the following
+three shared layers. The indexer implements the checkpoint equations in fp32:
+Q-LoRA residual projection, LayerNorm+interleaved-RoPE key projection, ReLU
+scores, learned signed head weights, causal bias, and dynamic
+`min(2048, key_length)` TopK.
+
+Each full-indexer layer packs its 128-value index key beside the decompressed
+MLA key in the standard `past_key_values.N.key` tensor. Shared layers retain
+the normal MLA cache. This keeps the ORT GenAI key/value naming contract
+without duplicating the index key across 64 attention heads.
+
+The MTP component exports:
+
+1. `enorm(next-token embedding)` and `hnorm(target hidden state)`;
+2. concatenation and the `2H -> H` `eh_proj`;
+3. the complete layer-78 DSA+MoE decoder block; and
+4. `shared_head.norm`, producing `mtp_hidden` for the shared target LM head.
+
+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.
+
+The importer:
+
+- derives all MLA, MoE, IndexShare, RoPE, and NextN fields from GGUF metadata;
+- separates the GGUF MTP block from the 78-layer backbone and exports it as `mtp`;
+- fuses split `attn_k_b` and `attn_v_b` tensors into `kv_b_proj`;
+- maps `blk.N.indexer.*` and supported `blk.78.nextn.*` tensors;
+- streams stacked routed-expert tensors one expert at a time instead of
+  materializing all 256 experts as one float tensor; and
+- preserves supported 4-bit and 8-bit tensors as MatMulNBits, requantizing
+  mixed source types to the selected graph-wide layout when necessary.
+
+The generic portable MoE implementation evaluates every expert and masks the
+outputs. It is correct for routing, but a fused sparse-MoE runtime kernel will
+be required for practical GLM-5.2 execution.
+
+## Runtime follow-ups
+
+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.
+
+`index_share_for_mtp_iteration` also needs runtime orchestration. MTP step 0 can
+return `topk_indices`, but accepted speculative positions make the MTP MLA and
+indexer caches advance at different logical lengths. ORT GenAI currently has
+no separate indexer-cache/control-state contract, so steps 1+ recompute the
+index selection rather than silently applying an incorrect packed-cache reuse.
+
+Use `--glm-full-attention` to retain the prior dense MLA fallback. The fallback
+disables the DSA-dependent MTP artifact.
+
+## Separate workstream: Unsloth dynamic quantization
+
+The six `UD-IQ1_M` GGUF shard headers were inspected. The dominant routed
+expert tensors are below the currently targeted 4-bit floor:
+
+- routed gate/up projections: 76 `IQ1_M` and 74 `IQ2_XXS` tensors;
+- routed down projections: 71 `IQ3_XXS` and 4 `IQ4_XS` tensors;
+- shared experts: mainly `Q5_K` / `Q6_K`;
+- MLA projections: mainly `Q8_0` / `Q5_K`;
+- norms: F32; embedding and output: Q4_K.
+
+This increment can dequantize IQ sources and requantize them to a supported
+Q4 MatMulNBits layout. That enables export but loses the intended sub-4-bit
+size and may alter accuracy. Native preservation should use one or both of the
+following paths.
+
+### Runtime custom-op requirement
+
+Add an ONNX Runtime GenAI execution-provider capability for a fused sparse-MoE
+operator that:
+
+1. consumes top-k expert indices and routing weights without evaluating all
+   256 experts;
+2. supports the llama.cpp IQ1_M, IQ2_XXS, and IQ3_XXS block/codebook layouts
+   (plus IQ4_XS), including their scales/minimums;
+3. fuses gate/up activation and down projection where profitable;
+4. loads only selected expert blocks and supports CPU and target accelerator
+   implementations; and
+5. advertises supported bit layouts through Mobius EP capabilities so export
+   selects native packed initializers, otherwise falling back to Q4
+   MatMulNBits requantization.
+
+### Draft ONNX Runtime issue
+
+**Title:** Support sub-4-bit GGUF/IQ weights for MatMulNBits and sparse MoE
+
+**Body:**
+
+> GLM-5.2 dynamic GGUF distributions use IQ1_M and IQ2_XXS for most routed
+> gate/up expert matrices and IQ3_XXS for most down matrices. Exporting these
+> models to ONNX currently requires expansion or requantization to 4-bit,
+> losing the distribution's size advantage and making a 256-expert model
+> impractical. Please define supported 1/2/3-bit MatMulNBits representations,
+> or provide an extensible packed/codebook-weight contract for the llama.cpp
+> IQ formats, with CPU and execution-provider kernels. For MoE, an operator
+> that accepts expert weights plus per-token top-k indices/weights should avoid
+> evaluating or unpacking every expert. Required formats for the first target
+> are IQ1_M, IQ2_XXS, IQ3_XXS, and IQ4_XS. Please also document zero-point,
+> scale/codebook packing, block-size, alignment, and external-data constraints.

export_deepseek_v2_tiny.py

@@ -0,0 +1,74 @@
+"""Export the tiny synthetic DeepSeek-V2 MLA + MoE model for onnx-genai E2E."""
+
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "tests"))
+
+import numpy as np
+import onnx_ir as ir
+from _test_configs import ALL_CAUSAL_LM_CONFIGS, _base_config
+
+from mobius._config_resolver import _default_task_for_model
+from mobius._registry import registry
+from mobius.integrations.onnx_genai import write_onnx_genai_config
+from mobius.tasks import get_task
+
+ARTIFACTS_DIR = os.environ.get(
+    "MOBIUS_ARTIFACTS_DIR", os.path.join(os.path.dirname(__file__), "artifacts")
+)
+
+
+def _fill_random_weights(model: ir.Model, rng: np.random.Generator) -> None:
+    for init in model.graph.initializers.values():
+        if init.const_value is not None:
+            continue
+        shape = tuple(int(d) for d in init.shape)
+        if not shape:
+            continue
+        if init.dtype == ir.DataType.FLOAT:
+            data = rng.standard_normal(shape).astype(np.float32) * 0.02
+        elif init.dtype == ir.DataType.FLOAT16:
+            data = (rng.standard_normal(shape) * 0.02).astype(np.float16)
+        elif init.dtype in (ir.DataType.INT64, ir.DataType.INT32):
+            npd = np.int64 if init.dtype == ir.DataType.INT64 else np.int32
+            data = rng.integers(0, 10, size=shape).astype(npd)
+        else:
+            data = rng.standard_normal(shape).astype(np.float32) * 0.02
+        init.const_value = ir.Tensor(data)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--output-dir",
+        default=os.path.join(ARTIFACTS_DIR, "deepseek-v2-tiny"),
+    )
+    args = parser.parse_args()
+
+    overrides = dict(next(ov for mt, ov, _ in ALL_CAUSAL_LM_CONFIGS if mt == "deepseek_v2"))
+    config = _base_config(**overrides)
+    config.dtype = ir.DataType.FLOAT
+
+    module = registry.get("deepseek_v2")(config)
+    task = get_task(_default_task_for_model("deepseek_v2"))
+    pkg = task.build(module, config)
+
+    rng = np.random.default_rng(0)
+    for model in pkg.values():
+        _fill_random_weights(model, rng)
+
+    os.makedirs(args.output_dir, exist_ok=True)
+    pkg.save(args.output_dir, external_data="onnx", check_weights=False)
+    artifacts = write_onnx_genai_config(pkg, args.output_dir, config=config)
+    for name, path in artifacts.items():
+        print(f"{name}:", path)
+    print("Saved to", args.output_dir)
+    print("files:", sorted(os.listdir(args.output_dir)))
+
+
+if __name__ == "__main__":
+    main()

export_glm_tiny_qmoe.py

@@ -0,0 +1,133 @@
+"""Export a tiny synthetic glm_moe_dsa model with a FUSED ``com.microsoft::QMoE``.
+
+Same tiny ``glm_moe_dsa`` config as ``export_glm_tiny_quant.py``, but sets
+``config.fused_quantized_moe = True`` so the routed MoE experts are emitted as a
+single fused ``com.microsoft::QMoE`` node (int4, block-32, expert-major layout)
+instead of a per-expert unroll of ``com.microsoft::MatMulNBits``.
+
+Routing (GLM sigmoid + noaux_tc, n_group=1) is preserved exactly: the QMoE
+kernel re-derives top-k selection from ``router_probs`` (the gate's
+``scores_for_choice``), and the normalized+scaled combine weights are scattered
+into the optional ``router_weights`` input with ``normalize_routing_weights=0``.
+
+Writes an onnx-genai-loadable artifact directory:
+  <out>/model.onnx (+ external data), inference_metadata.yaml, tokenizer.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import shutil
+import sys
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "tests"))
+
+import numpy as np
+import onnx_ir as ir
+from _test_configs import ALL_CAUSAL_LM_CONFIGS, _base_config
+
+from mobius._config_resolver import _default_task_for_model
+from mobius._configs import QuantizationConfig
+from mobius._registry import registry
+from mobius.integrations.onnx_genai import write_onnx_genai_config
+from mobius.tasks import get_task
+
+ARTIFACTS_DIR = os.environ.get(
+    "MOBIUS_ARTIFACTS_DIR", os.path.join(os.path.dirname(__file__), "artifacts")
+)
+
+
+def _fill_random_weights(model: ir.Model, rng: np.random.Generator) -> None:
+    for init in model.graph.initializers.values():
+        if init.const_value is not None:
+            continue
+        shape = tuple(int(d) for d in init.shape)
+        if not shape:
+            continue
+        if init.dtype == ir.DataType.FLOAT:
+            data = rng.standard_normal(shape).astype(np.float32) * 0.02
+        elif init.dtype == ir.DataType.FLOAT16:
+            data = (rng.standard_normal(shape) * 0.02).astype(np.float16)
+        elif init.dtype == ir.DataType.UINT8:
+            # Packed int4/int8 quantized weights or bit-packed zero points.
+            data = rng.integers(0, 256, size=shape).astype(np.uint8)
+        elif init.dtype in (ir.DataType.INT64, ir.DataType.INT32):
+            npd = np.int64 if init.dtype == ir.DataType.INT64 else np.int32
+            data = rng.integers(0, 10, size=shape).astype(npd)
+        else:
+            data = rng.standard_normal(shape).astype(np.float32) * 0.02
+        init.const_value = ir.Tensor(data)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--output-dir",
+        default=os.path.join(ARTIFACTS_DIR, "glm-5.2-tiny-qmoe"),
+    )
+    parser.add_argument(
+        "--fp32-output-dir",
+        default=os.path.join(ARTIFACTS_DIR, "glm-5.2-tiny"),
+    )
+    args = parser.parse_args()
+
+    overrides = dict(next(ov for mt, ov, _ in ALL_CAUSAL_LM_CONFIGS if mt == "glm_moe_dsa"))
+    config = _base_config(**overrides)
+    config.dtype = ir.DataType.FLOAT
+    # int4, block-32, symmetric -> fused QMoE (no zero points; the kernel
+    # defaults per-block zero-point to 1 << (bits - 1)).
+    config.quantization = QuantizationConfig(
+        bits=4,
+        group_size=32,
+        quant_method="gguf",
+        sym=True,
+    )
+    # Emit routed experts as a single fused com.microsoft::QMoE node.
+    config.fused_quantized_moe = True
+
+    model_type = "glm_moe_dsa"
+    module = registry.get(model_type)(config)
+    task = get_task(_default_task_for_model(model_type))
+    pkg = task.build(module, config)
+
+    rng = np.random.default_rng(0)
+    for model in pkg.values():
+        _fill_random_weights(model, rng)
+
+    os.makedirs(args.output_dir, exist_ok=True)
+    pkg.save(args.output_dir, external_data="onnx", check_weights=False)
+    artifacts = write_onnx_genai_config(pkg, args.output_dir, config=config)
+    for name, path in artifacts.items():
+        print(f"{name}:", path)
+
+    # Reuse tokenizer from the fp32 artifacts if present.
+    tok_src = os.path.join(args.fp32_output_dir, "tokenizer.json")
+    if os.path.exists(tok_src):
+        shutil.copy(tok_src, os.path.join(args.output_dir, "tokenizer.json"))
+
+    for attr in (
+        "num_hidden_layers",
+        "num_attention_heads",
+        "num_key_value_heads",
+        "head_dim",
+        "hidden_size",
+        "vocab_size",
+        "max_position_embeddings",
+        "qk_nope_head_dim",
+        "qk_rope_head_dim",
+        "v_head_dim",
+        "kv_lora_rank",
+        "q_lora_rank",
+        "num_local_experts",
+        "num_experts_per_tok",
+        "moe_intermediate_size",
+    ):
+        print(f"  {attr} =", getattr(config, attr, None))
+
+    print("Saved to", args.output_dir)
+    print("files:", sorted(os.listdir(args.output_dir)))
+
+
+if __name__ == "__main__":
+    main()