Qwen3.6: export ModelOpt NVFP4/FP8 checkpoints natively by tianleiwu · Pull Request #2351 · microsoft/onnxruntime-genai

Tianlei Wu (tianleiwu) · 2026-07-31T09:34:32Z

Description

Stacked on #2350 (tlwu/20260731/nvfp4_quant_core). Review only the last commit; the base branch will be retargeted to main once #2350 merges.

Wires the NVFP4 / FP8 scheme into the Qwen3.6 MoE builder so a TensorRT Model Optimizer quantized checkpoint is carried into the ONNX model instead of being dequantized to fp16 and re-quantized to int4.

NVFP4 routed experts — moe_quant_type=nvfp4 emits a native NVFP4 QMoE op built straight from the checkpoint's packed E2M1 codes and E4M3 block scales (block size 16).
NVFP4 dense modules — use_original_nvfp4_weights=true emits the shared expert and lm_head as the weight-only MatMulBlockQuantizedFp4Weight contrib op, feeding the checkpoint tensors through unmodified. nvfp4_dense_exclude_layers / nvfp4_lmhead_fp16 keep selected modules at fp16 for A/B.
FP8 dense projections — use_original_fp8_weights=true emits the self-attention q/k/v/o projections as MatMulBlockQuantizedFp8Weight; fp8_linear_attn (default on) does the same for the GatedDeltaNet in_proj_qkv / in_proj_z / out_proj. Optional calibrated per-tensor activation scales via fp8_attn_static_input_scale, fp8_linear_attn_static_input_scale and share_fp8_attn_qkv_activation; fp8_attn_exclude_layers isolates a layer for error analysis.
FP8 KV cache — fp8_kv_cache=true with kv_cache_scale_file=... stores the KV cache as E4M3, supporting both calibrated per-layer scales (see Support hybrid (linear-attention) models in KV cache calibration #2348) and the legacy shared unit scale.
Gate fusion — fuse_linear_attn_gates (default on for CUDA) collapses the float32 gate glue around LinearAttention into com.microsoft::LinearAttentionGate and com.microsoft::GatedRMSNorm, removing ~9 nodes per linear-attention layer and the corresponding CUDA-graph replay overhead.

Motivation and Context

The reference Qwen3.6 checkpoint is already NVFP4/FP8 quantized. Round-tripping through fp16 loses the calibrated scales and inflates both the export and the runtime weight traffic. These options require an ONNX Runtime build providing the corresponding contrib ops.

Adds unit tests for the FP8 activation-scale and FP8 linear-attention paths, and documents the new extra options in the model builder README.

Adds `nvfp4` alongside `mxfp4` as a MoE quantization scheme. NVFP4 uses 4-bit e2m1 weights with FP8-E4M3 block scales over blocks of 16 (vs ue8m0 scales over blocks of 32 for MXFP4) plus a per-expert float32 global scale, and selects the CUDA QMoE `quant_type="nvfp4"` kernel. Both schemes share the FP4 QMoE input layout (global scales at input positions 15/16) and the same build requirements: CUDA EP and symmetric int4 build precision. - quant_config.py: `nvfp4` dtype descriptor pinned to block size 16. - base.py: map the MoE dtype name to the QMoE `quant_type` attribute, emit the FP4 global-scale inputs for both schemes, and add `make_fp8e4m3_initializer` for FP8-E4M3 block scales plus `repack_modelopt_nvfp4_weight_codes` / `pack_nvfp4_codes_for_qmoe` to convert a Model Optimizer `[N, K/2]` weight tensor into the QMoE `[K, N/2]` layout. - builder.py: accept and document `moe_quant_type=nvfp4`. - quantized_model.py: `ModeloptModel`, a loader for NVIDIA Model Optimizer mixed-precision checkpoints (NVFP4 experts / shared expert / lm_head, FP8 attention, BF16 everything else). It dequantizes the non-routed weights to BF16 for the builder's normal path; the routed experts are streamed from the source safetensors by the model builder instead.

Carry a TensorRT Model Optimizer quantized Qwen3.6 MoE checkpoint into the ONNX model instead of dequantizing to fp16 and re-quantizing to int4. Builder changes (src/python/py/models/builders/qwen.py): * Routed MoE experts are emitted as a native NVFP4 QMoE op when moe_quant_type=nvfp4 (E2M1 codes + E4M3 block scales, block size 16), loading the packed checkpoint tensors directly. * Dense NVFP4 modules (shared expert, lm_head) are emitted as the weight-only MatMulBlockQuantizedFp4Weight contrib op under use_original_nvfp4_weights, with nvfp4_dense_exclude_layers / nvfp4_lmhead_fp16 escape hatches. * Self-attention q/k/v/o and the GatedDeltaNet in_proj_qkv / in_proj_z / out_proj projections are emitted as MatMulBlockQuantizedFp8Weight under use_original_fp8_weights / fp8_linear_attn, with optional calibrated per-tensor activation scales (fp8_attn_static_input_scale, fp8_linear_attn_static_input_scale, share_fp8_attn_qkv_activation) and fp8_attn_exclude_layers for A/B isolation. * FP8 (E4M3) KV cache via fp8_kv_cache + kv_cache_scale_file, supporting both calibrated per-layer scales and the legacy shared unit scale. * fuse_linear_attn_gates (default on for CUDA) collapses the float32 gate glue around LinearAttention into com.microsoft::LinearAttentionGate and com.microsoft::GatedRMSNorm, removing ~9 nodes per linear-attention layer. Adds unit tests for the FP8 activation-scale and FP8 linear-attention paths, and documents the new extra options in the model builder README.

Copilot

Pull request overview

Adds native export support in the Qwen3.6 MoE builder for Model Optimizer–quantized checkpoints by preserving original NVFP4/FP8 tensors in the ONNX graph (instead of dequantizing to fp16 and re-quantizing), plus documentation and focused unit tests for the new FP8 behaviors.

Changes:

Extend Qwen35MoeTextModel to emit weight-only FP8 (MatMulBlockQuantizedFp8Weight) for attention and (optionally) linear-attention projections, and weight-only NVFP4 (MatMulBlockQuantizedFp4Weight) for dense NVFP4 modules.
Add hybrid-stack FP8 KV-cache handling, including a legacy compatibility mode and per-KV-layer scale initialization support.
Add unit tests for FP8 activation-scale initializer sharing and for linear-attention FP8 checkpoint key mapping; document the new Qwen3.6 ModelOpt export flow in the builder README.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 3 comments.

File	Description
test/python/builder/test_qwen_fp8_linear_attn.py	New unit tests validating FP8 checkpoint key mapping for linear-attention matmuls and opt-out behavior.
test/python/builder/test_qwen_fp8_activation.py	New unit tests validating static FP8 activation scale loading and initializer sharing behavior.
src/python/py/models/README.md	Documents the new ModelOpt NVFP4/FP8 export scenario and related extra options for Qwen3.6.
src/python/py/models/builders/qwen.py	Implements native NVFP4/FP8 export wiring for Qwen3.6, FP8 KV-cache handling, and optional linear-attention gate fusion.

Add an opt-in path that stores the com.microsoft::GroupQueryAttention past/present KV cache as FLOAT8E4M3FN instead of the model dtype, halving KV-cache memory at long context. Mirrors onnxruntime-genai's fp8_kv_cache option (microsoft/onnxruntime-genai#2351). A new post-fusion IR pass, Fp8KvCachePass, retypes each decoder GQA past_key/past_value input and present_key/present_value output to FP8, adds per-layer k_scale/v_scale FLOAT initializers at GQA input slots 12/13, and sets k_quant_type/v_quant_type="PER_TENSOR" and kv_cache_bit_width=8. The query/key/value inputs stay at the model dtype; the kernel quantizes the new K/V to FP8 on write and dequantizes on read using the scales. Scales default to a unit 1.0 (the "legacy" export shape) or come from an offline calibration file (onnxruntime-genai scales JSON format). Wired through optimize_model, build, build_from_module, and the `mobius build` CLI via --fp8-kv-cache and --kv-cache-scale-file. The option is ignored with a warning when GQA fusion is not active (e.g. non-CUDA EP or fp32), since there is no KV-cache op to convert. Verified on an H200 (SM90) with onnxruntime-gpu 1.28: the emitted FP8 GQA op loads and computes on the CUDA EP (finite fp16 output, FP8 present cache). End-to-end FP8 KV at runtime is IO-bound as device FLOAT8E4M3FN OrtValues (as onnxruntime-genai does); the Python numpy fp8 feed path is unsupported by ORT, which the runtime test works around with empty in-graph FP8 constants. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: justinchuby <justinchuby@users.noreply.github.com>

## Summary Adds an opt-in **FP8 (E4M3) KV-cache export path**. When enabled, each decoder `com.microsoft::GroupQueryAttention` stores its past/present key-value cache as `FLOAT8E4M3FN` instead of the model dtype (fp16/bf16), **halving KV-cache memory** at long context. Mirrors onnxruntime-genai's `fp8_kv_cache` option (microsoft/onnxruntime-genai#2351). ## How it works A new post-fusion IR pass `Fp8KvCachePass` (`src/mobius/_passes/_fp8_kv_cache.py`), run after the GQA fusion rules, for every decoder GQA node whose KV cache are graph inputs: 1. Retypes the `past_key`/`past_value` inputs and `present_key`/`present_value` outputs (all graph I/O) to `FLOAT8E4M3FN`, preserving shapes. 2. Adds per-layer `k_scale`/`v_scale` scalar FLOAT initializers at GQA input slots **12/13** (ORT ≥ 1.28 signature). 3. Sets `k_quant_type`/`v_quant_type = "PER_TENSOR"` and `kv_cache_bit_width = 8`. `query`/`key`/`value` stay at the model dtype — the kernel quantizes the new K/V to FP8 on write and dequantizes on read using the scales. Scales default to a unit `1.0` (the "legacy" export shape), or come from an offline **calibration file** (onnxruntime-genai `{"scales": {"k_scales": [...], "v_scales": [...]}}` JSON). ## Usage ```bash mobius build --model Qwen/Qwen3-4B --dtype f16 \ --execution-provider cuda --fp8-kv-cache \ [--kv-cache-scale-file scales.json] /out ``` Also exposed programmatically on `build(...)` / `build_from_module(...)` via `fp8_kv_cache=` and `kv_cache_scales=`. The option is **ignored with a warning** when GQA fusion is not active (non-GQA EP or fp32), since there is no KV-cache op to convert. ## Verification (H200 / SM90, onnxruntime-gpu 1.28) - The emitted FP8 GQA op **loads and computes on the CUDA EP** — finite fp16 output, FP8 present cache (covered by `test_fp8_kv_gqa_runs_on_cuda`, gated on CUDA availability). - End-to-end FP8 KV at runtime is IO-bound as device `FLOAT8E4M3FN` OrtValues (as onnxruntime-genai does). The Python numpy fp8 **feed** path is unsupported by ORT; the runtime test works around it with empty in-graph FP8 constants. ## Tests `src/mobius/_passes/_fp8_kv_cache_test.py` — 11 tests: KV I/O typed FP8, GQA scale inputs + quant attrs, default unit scales, calibrated scales, disabled path, ignored+warns on non-GQA EP, idempotency, scale-file parsing/validation, and the CUDA runtime run. Existing pass/CLI/build-graph suites pass; changed files pass ruff. > ⚠️ Please review — do not merge yet. --------- Signed-off-by: justinchuby <justinchuby@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>

…#446) ## What Foundation for consuming **NVIDIA ModelOpt** NVFP4/FP8 mixed-precision checkpoints (e.g. quantized **Qwen3.6**), porting the verifiable numeric core of ORT GenAI [#2350](microsoft/onnxruntime-genai#2350) / [#2351](microsoft/onnxruntime-genai#2351). ModelOpt Qwen3.6 layout: - **Routed MoE experts / shared expert / lm_head** → `W4A16_NVFP4`: block-16 E2M1 (fp4) weights, FP8-E4M3 block scales, per-tensor FP32 global scale (`weight_scale_2`). - **Attention / linear-attn projections** → `FP8` (E4M3, per-tensor `weight_scale`). ## How New `mobius/integrations/modelopt/` package with the loader's numeric core: - `dequantize_nvfp4()` — `e2m1(code) × e4m3(block_scale[k//16]) × global_scale → bf16` - `dequantize_fp8()` — `e4m3(weight) × weight_scale → bf16` - `unpack_nvfp4_codes()`, `is_modelopt_quant_config()`, `FP4_E2M1_LUT` Following ORT GenAI's loader, dense FP8 and NVFP4 shared-expert/lm_head weights are **reconstructed to BF16 exactly** (ORT has no FP8 attention GEMM), so the standard BF16 build path can consume the checkpoint. Fully unit-tested against hand-computed references — numeric-only, no runtime/GPU dependency. `QuantizationConfig.from_transformers` now **fails loudly** (`NotImplementedError`) on ModelOpt schemes instead of silently routing packed E2M1/float8 weights through the INT4 `MatMulNBits` path (which would mis-dequantize them). The check runs before the `quant_method=="none"` early-return since ModelOpt may name its scheme only via `quant_algo`/`quant_cfg`. ## Scope / verification - ✅ Verifiable here: dequant math (7 unit tests, hand-computed references) + config detection/guard (89 tests pass in `_configs_test.py` + module). - ⏳ **Out of scope (documented follow-up):** full checkpoint weight-load wiring and **native routed-expert NVFP4 QMoE emission** (CUDA/Blackwell, `onnxruntime_USE_FP4_QMOE=ON`). That path is **not buildable or verifiable** in this environment (no Blackwell GPU, no FP4-QMOE ORT build, no real ModelOpt checkpoint), so it is intentionally gated rather than shipped unverified — the safe increment given the repo's "no silently-wrong weights" bar. Please review — do not merge yet. --------- Signed-off-by: justinchuby <justinchuby@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Code scanning does not honor lgtm[...] suppression comments, so the py/overwritten-inherited-attribute alert kept firing. Restoring the original lines also takes this pre-existing code back out of the PR diff.

…nvfp4_fp8_builder

Tianlei Wu (tianleiwu) · 2026-08-14T20:34:15Z

/azp run Integration Tests

azure-pipelines · 2026-08-14T20:34:27Z

Azure Pipelines: Successfully started running 1 pipeline(s).

## Summary Adds Qwen3.6 MoE multi-token-prediction (MTP) head export for the self-speculative decoding runtime merged in #2352. With `--extra_options enable_mtp=true include_hidden_states=true`, the builder writes `mtp.onnx` alongside the main decoder and adds the corresponding `model.mtp` contract to `genai_config.json`. This completes the builder side of Qwen3.6 MTP while preserving native ModelOpt NVFP4/FP8 tensors by default and allowing the MTP model to use an independent structured quantization configuration. ## Key Changes ### MTP head export - Adds `Qwen35MtpHead`, a single full-attention Qwen3.5 MoE decoder layer that predicts the next-next token from: - the main decoder's last hidden state; and - the embedding of the just-emitted token. - Reads `mtp.*` tensors directly from source safetensors because Hugging Face `transformers` does not retain them in the loaded model. - Exports `hidden_states_out` for recurrent multi-token drafting. - Adds `model.mtp` filename, dimensions, inputs, outputs, and main hidden-state mapping to `genai_config.json`. - Validates that MTP export keeps main-model hidden states and the LM head enabled. ### Independent MTP quantization configuration - By default, the MTP model inherits the main model's settings. - ModelOpt checkpoints preserve each MTP tensor's native format by default: - NVFP4 linears and routed experts remain NVFP4; - FP8 projections remain FP8; and - unquantized tensors use the graph precision. - Adds `mtp_quant_config`, which accepts inline JSON or a JSON file using the structured `QuantConfig` schema. - Resolves MTP I/O dtype, dense MatMul weights, MoE experts, runtime settings, and mixed-precision overrides independently from the main model. - An explicit MTP configuration dequantizes native ModelOpt MTP tensors before applying the requested formats. - Supports unquantized or INT4/INT8 dense MTP weights and independently configured INT4/INT8/MXFP4/NVFP4 MoE experts. - Supports independent calibrated MTP KV-cache scales through the optional `mtp` section of `kv_cache_scale_file`. ### External-data deduplication - Shares byte-identical embedding and `lm_head.MatMul.*` initializers between the main model and MTP graph. - Repacks `mtp.onnx.data` after redirected tensors are removed. - Stages both ONNX metadata and external data before replacing either file, and detects truncated input data so a failed optimization leaves the original export intact. ### Windowed recurrent state - Adds `state_window=W` to widen Qwen3.6 convolution and recurrent state I/O to `[W, B, ...]`. - Emits the matching `state_window` attribute on `CausalConvWithState` and `LinearAttention`. - Allows a multi-token verification forward to crop state to the accepted prefix instead of replaying the decoder. - Requires `W >= num_speculative_tokens + 1` for MTP verification. The default `0` retains legacy unwindowed I/O. ### Builder robustness - Guards cleanup of the cache directory shared by the main and MTP exports. - Validates ModelOpt tensor metadata and explicit quantization selections. - Resolves Qwen3.5 dense, MoE, and MTP runtime model types during base initialization rather than overwriting inherited attributes. - Documents the supported MTP export options and windowed-state behavior. ## Usage ```bash python -m onnxruntime_genai.models.builder \ -i <qwen3.6-model-directory> \ -o <output-directory> \ -p <precision> \ -e <execution-provider> \ --extra_options \ enable_mtp=true \ include_hidden_states=true \ state_window=4 ``` Optional independent MTP quantization configuration: ```bash --extra_options \ enable_mtp=true \ include_hidden_states=true \ mtp_quant_config='{"io_dtype":"bf16","weights":{"type":"int4","block_size":64},"moe":{"type":"nvfp4"}}' ``` To keep the complete MTP model in FP16, use: ```text mtp_quant_config='{"io_dtype":"fp16","weights":{"type":"none"},"moe":{"type":"none"}}' ``` MTP export is currently supported for Qwen3.6 MoE checkpoints with architecture `Qwen3_5MoeForConditionalGeneration` and source safetensors containing `mtp.*` weights. ## Validation - `PYTHONPATH=src/python/py python -m pytest -q test/python/builder/test_mixed_precision_config.py test/python/builder/test_quant_config.py test/python/builder/test_qwen_mtp_head_quant.py test/python/builder/test_qwen_mtp_export.py test/python/builder/test_precision.py` - `124 passed` - `lintrunner` passed for the touched Python builder and test files. - `git diff --check` passed. - Qwen3.6-35B-A3B ModelOpt NVFP4 MTP export completed end to end before the configuration surface was generalized. - Full PR CI is in progress. ## Related PRs - #2350: Model builder NVFP4 quantization core, merged. - #2351: Qwen3.6 native ModelOpt NVFP4/FP8 export, merged. - #2352: MTP self-speculative decoding runtime, merged. - Supersedes the closed #2218. --------- Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com>

## Description Fuse Qwen3.6 MoE shared-expert scaling and routed/shared addition into the `com.microsoft::GatedAdd` contrib operator. Qwen3.6 shared-expert Mul+Add is replaced by com.microsoft::GatedAdd for CPU, CUDA, and WebGPU. Unsupported execution providers retain portable ONNX `Mul` + `Add`. No `fuse_shared_expert_gate` option remains. This PR is stacked on #2353, which is itself stacked on #2351. Dependencies are microsoft/onnxruntime#31835 for CUDA and merged microsoft/onnxruntime#32106 for CPU/WebGPU. ## Changes - Return the shared-expert projection and scalar gate separately from `make_shared_expert`. - Emit one `GatedAdd` per MoE layer on CPU, CUDA, and WebGPU. - Preserve the standard `Mul` + `Add` graph as an explicit fallback for unsupported EPs. - Add focused tests for fused and fallback graph construction. ## Performance On Qwen3.6-35B-A3B-NVFP4 with N=3 MTP, the real exported graph replaces 40 main-model pairs plus one MTP pair. Counterbalanced H200 measurements reduced median decode latency from 7.311 to 7.225 ms/round (-1.18%). Graph-off Nsight measured 40.35 fewer launches/round and 1.30% lower GPU kernel time. ## Validation - `50 passed`: `test_precision.py` plus `test_qwen_gated_add.py`. - Real graph census: 40 `GatedAdd` nodes in `text.onnx`, one in `mtp.onnx`. - After normalizing fused edge names, all other nodes, initializers, graph inputs, and graph outputs are unchanged from the baseline export. - Runtime float, FP16, and BF16 results are bit-exact with separate `Mul` + `Add`. ## Stack note The source commit `5c35bb02fc` also contained an unrelated MTP prefill chunk default change in `src/mtp_generator.{h,cpp}`. Those files are intentionally excluded here because they depend on runtime PR #2352 rather than builder PR #2353.

Tianlei Wu (tianleiwu) added 2 commits July 31, 2026 09:30

Tianlei Wu (tianleiwu) requested a review from a team as a code owner July 31, 2026 09:34

Tianlei Wu (tianleiwu) mentioned this pull request Jul 31, 2026

Qwen3.6: build the MTP self-speculative head #2353

Merged

Tianlei Wu (tianleiwu) marked this pull request as draft July 31, 2026 17:51

Justin Chu (justinchuby) mentioned this pull request Aug 1, 2026

Add NVFP4/FP8 (ModelOpt) weight reconstruction foundation for Qwen3.6 onnxruntime/mobius#446

Merged

Tianlei Wu (tianleiwu) requested a review from Copilot August 1, 2026 20:44

Copilot started reviewing on behalf of Tianlei Wu (tianleiwu) August 1, 2026 20:45 View session

Copilot AI reviewed Aug 1, 2026

View reviewed changes

Comment thread src/python/py/models/builders/qwen.py Outdated

Comment thread src/python/py/models/builders/qwen.py Outdated

Comment thread src/python/py/models/builders/qwen.py Outdated

Justin Chu (justinchuby) mentioned this pull request Aug 2, 2026

feat(kv-cache): add FP8 (E4M3) KV-cache export path onnxruntime/mobius#447

Merged

Base automatically changed from tlwu/20260731/nvfp4_quant_core to main August 7, 2026 20:12

Tianlei Wu (tianleiwu) added 3 commits August 8, 2026 06:37

Merge origin/main

27f6ff5

scales shape and add tests

437748e

refine

8148347

Tianlei Wu (tianleiwu) marked this pull request as ready for review August 8, 2026 07:13

github-advanced-security Bot found potential problems Aug 8, 2026

View reviewed changes

Comment thread src/python/py/models/builders/qwen.py Fixed

Comment thread src/python/py/models/builders/qwen.py Fixed

remove experimental options; address feedbacks.

3fb45b0

Tianlei Wu (tianleiwu) mentioned this pull request Aug 10, 2026

Qwen3.6: fuse shared-expert gating with GatedAdd #2390

Merged

Tianlei Wu (tianleiwu) requested a review from kunal-vaishnavi August 10, 2026 03:42

Tianlei Wu (tianleiwu) added the model:qwen-3.6 label Aug 11, 2026

fix(qwen): handle optional FP8 input scales

609e98f

github-advanced-security Bot found potential problems Aug 11, 2026

View reviewed changes

Comment thread src/python/py/models/builders/qwen.py Fixed

Tianlei Wu (tianleiwu) added 2 commits August 11, 2026 06:03

fix(qwen): drop ineffective CodeQL suppression comments

a18a870

Code scanning does not honor lgtm[...] suppression comments, so the py/overwritten-inherited-attribute alert kept firing. Restoring the original lines also takes this pre-existing code back out of the PR diff.

Merge remote-tracking branch 'origin/main' into tlwu/20260731/qwen36_…

19170a4

…nvfp4_fp8_builder