Add native GPT-OSS MXFP4 QMoE export - #719
Open
titaiwangms wants to merge 7 commits into
Open
Conversation
Preserve official GPT-OSS MXFP4 expert weights during QMoE export and stream indexed safetensors one projection at a time. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 1d9a8b1f-9b83-4a52-bfe4-92437196e1c1
Performance Comparison
|
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Apply the existing streaming save policy and revalidate E8M0 scales when lazy sources are materialized. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 1d9a8b1f-9b83-4a52-bfe4-92437196e1c1
Emit the required ONNX IR version and prevent streaming safetensors output from overwriting lazy source shards through directory, symlink, component-layout, or hard-link aliases. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 1d9a8b1f-9b83-4a52-bfe4-92437196e1c1
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
justinchuby
reviewed
Sep 3, 2026
justinchuby
reviewed
Sep 3, 2026
justinchuby
reviewed
Sep 3, 2026
Align native MXFP4 attention and KV-cache behavior with the ORT GenAI model builder, and keep graph capture consistent with the effective shared-buffer contract. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 1d9a8b1f-9b83-4a52-bfe4-92437196e1c1
Contributor
There was a problem hiding this comment.
🔵 Needs a closer look
It modifies core export/weight-loading/serialization paths and introduces a new native FP4 runtime-dependent execution slice that needs final qualification on official checkpoints and the targeted ORT CUDA build.
Pull request overview
This PR adds a native export path for GPT-OSS MXFP4 MoE checkpoints that preserves the original FP4 expert codes/scales (no dequant/requant), repacks them into ORT’s native com.microsoft::QMoE FP4 layout, and introduces a bounded-memory streaming safetensors loader so very large checkpoints (e.g., 120B) can be exported without holding an eager second copy of expert weights in host RAM.
Changes:
- Implement native GPT-OSS MXFP4 export: lossless FP4 nibble repacking, E8M0 byte validation, native QMoE emission, and direct
GroupQueryAttentionpath for the supported CUDA slice. - Add native MXFP4 safetensors streaming: new transformed streaming sources, plan validation, and
ModelPackage.save()collision guards to prevent overwriting lazily-read source shards. - Make export/config more robust: infer minimum ONNX IR version from graph dtypes (IR12 for FLOAT8E8M0), introduce
QuantizedWeightFormat, and resolve ORT GenAI graph-capture settings based on the finalized KV-cache contract.
File summaries
| File | Description |
|---|---|
| src/mobius/tasks/_task_test.py | Adds unit coverage for _make_model() IR version floor behavior (IR11 vs IR12). |
| src/mobius/tasks/_base.py | Sets model ir_version via minimum_ir_version(graph) instead of hard-coding 11. |
| src/mobius/models/gptoss.py | Implements native MXFP4 QMoE export, repacking/validation, and direct GQA path for supported builds. |
| src/mobius/models/gptoss_test.py | Expands unit tests for repacking, native MXFP4 preprocessing/graph emission, and CUDA GQA parity checks. |
| src/mobius/models/base.py | Gates quantized-linear factory selection on QuantizedWeightFormat.INTEGER_AFFINE. |
| src/mobius/integrations/transformers/_gptoss_weights.py | Adds bounded-memory native MXFP4 safetensors streaming planner + binder. |
| src/mobius/integrations/transformers/_gptoss_weights_test.py | Adds extensive tests for streaming transforms, save collision guards, and roundtrip portability. |
| src/mobius/integrations/transformers/_builder.py | Validates native GPT-OSS MXFP4 export contract early and routes weight loading to streaming. |
| src/mobius/integrations/transformers/_builder_test.py | Tests streaming selection and early contract failures (EP/dtype). |
| src/mobius/integrations/ort_genai/genai_config.py | Resolves decoder graph capture only after KV-cache/search settings are finalized. |
| src/mobius/integrations/ort_genai/genai_config_test.py | Adds coverage for graph-capture resolution vs share-buffer and beam settings. |
| src/mobius/integrations/ort_genai/ep_config.py | Ensures DML graph capture can be explicitly disabled (emit option key). |
| src/mobius/integrations/ort_genai/ep_config_test.py | Adds tests for DML graph-capture default vs explicit disable behavior. |
| src/mobius/integrations/ort_genai/auto_export_test.py | Updates/extends tests to ensure capture tracks KV-cache contract. |
| src/mobius/integrations/_weight_loading.py | Adds StreamingTransformedWeightSource + synthesized targets support, and threads it through streaming validation/binding. |
| src/mobius/components/_moe.py | Refactors QMoE gate realization into a reusable helper for routing adapters. |
| src/mobius/_weight_utils.py | Prevents integer-affine QMoE logic from claiming MXFP4 configs. |
| src/mobius/_model_package.py | Adds streaming-save collision protection, freezes component filters, and adds E8M0 safetensors compatibility shim + lock. |
| src/mobius/_model_package_test.py | Adds tests for E8M0 safetensors save compatibility, predicate freezing, and streaming collision checks. |
| src/mobius/_ir_utils.py | Introduces dtype inspection helpers and minimum_ir_version() (IR12 floor for FLOAT8E8M0). |
| src/mobius/_configs/_quantization.py | Adds QuantizedWeightFormat enum and weight-format normalization + MXFP4 parsing. |
| src/mobius/_configs/init.py | Re-exports QuantizedWeightFormat. |
| src/mobius/_configs_test.py | Adds tests for weight-format defaults, normalization, and transformers MXFP4 parsing. |
Review details
- Files reviewed: 23/23 changed files
- Comments generated: 1
- Review effort level: Lite
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
ONNX saves can corrupt hard-linked lazy sources, and native QMoE execution lacks automated parity coverage.
Once you've addressed the issues Copilot identified, you can request another Copilot review.
Review details
Suppressed comments (1)
src/mobius/models/gptoss_test.py:106
- This test helper introduces explicit protobuf construction through
onnx.helper/onnx.TensorProto, which this repository prohibits even in tests. Build the parity graph withonnx_ir(ir.Value,ir.Node,ir.Graph, andir.Model) and serialize that model instead.
tensor = onnx.helper.make_tensor_value_info
inputs = [
tensor("query", onnx.TensorProto.FLOAT16, [batch, query_length, 16]),
tensor("key", onnx.TensorProto.FLOAT16, [batch, query_length, 8]),
tensor("value", onnx.TensorProto.FLOAT16, [batch, query_length, 8]),
- Files reviewed: 23/23 changed files
- Comments generated: 2
- Review effort level: Balanced
Resolve the latest main conflicts while preserving native GPT-OSS MXFP4 streaming and GQA support. Adopt universal IR 12, harden source-safe transactional serialization, and restore explicit dense dequantization across CLI entry points. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 1d9a8b1f-9b83-4a52-bfe4-92437196e1c1
Execute the emitted CUDA QMoE with production-repacked checkpoint tensors and compare routing, clipped SwiGLU, and bias semantics against an independent decoder for prefill and decode-shaped inputs. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 1d9a8b1f-9b83-4a52-bfe4-92437196e1c1
Test the shared build boundary after local config exports stopped using the removed private config and module-building helpers. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 1d9a8b1f-9b83-4a52-bfe4-92437196e1c1
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
[E, N, K/32, 16]blocks into the native QMoE[E, K, N/2]layoutcom.microsoft::QMoEwith GPT-OSS routing, expert biases, global scales, and SwiGLU semanticsRuntime requirements
Native GPT-OSS MXFP4 export currently requires:
Unsupported execution providers and dtypes fail before graph construction or weight I/O.
Draft validation remaining
The implementation blockers are resolved. Before marking this PR ready, the
remaining work is runtime qualification with official checkpoints:
memory, temporary disk use, output size, and wall time
Runtime build
including numerical parity and peak GPU memory
Validation
0x00through0xfeare preserved and0xffis rejected