Add 2-bit (uint2) Gemma 4 MoE support to the model builder

src/python/py/models/README.md

@@ -140,6 +140,18 @@ python -m onnxruntime_genai.models.builder -i path_to_local_folder_on_disk -o pa
 python builder.py -i path_to_local_folder_on_disk -o path_to_output_folder -p int4 -e execution_provider -c cache_dir_to_store_temp_files
 ```
 
+#### 2-bit (INT2) pre-quantized MoE
+
+For a model whose MoE experts have already been quantized to 2-bit with [AMD Quark](https://quark.docs.amd.com/) (group-wise `uint2`), pass `-p int2`. The builder consumes the pre-quantized expert weights, scales, and zero-points directly (no float re-quantization) and emits a fused `QMoE` op with `expert_weight_bits=2`. This path is currently exercised by the Gemma 4 MoE (`gemma-4-26B-A4B-it`) text model on the CPU execution provider.
+
+```bash
+# From wheel:
+python -m onnxruntime_genai.models.builder -i path_to_local_folder_on_disk -o path_to_output_folder -p int2 -e cpu -c cache_dir_to_store_temp_files
+
+# From source:
+python builder.py -i path_to_local_folder_on_disk -o path_to_output_folder -p int2 -e cpu -c cache_dir_to_store_temp_files
+```
+
 ### GGUF Model
 
 This scenario is where your float16/float32 GGUF model is already on disk.

src/python/py/models/builder.py

@@ -22,6 +22,8 @@
     ErnieModel,
     Gemma2Model,
     Gemma3Model,
+    Gemma4MoEModel,
+    Gemma4Model,
     GemmaModel,
     GPTOSSModel,
     GraniteModel,
@@ -242,7 +244,7 @@ def check_extra_options(
 
     # `moe_quant_type` is the single option that selects the MoE quantization scheme. It replaces the
     # older per-type flags (`use_8bits_moe``) so new schemes can be added without a new flag.
-    supported_moe_quant_types = {"int4", "int8", "mxfp4", "nvfp4"}
+    supported_moe_quant_types = {"int2", "uint2", "int4", "int8", "mxfp4", "nvfp4"}
 
     # Backward compatibility: `use_8bits_moe` is deprecated in favor of `moe_quant_type`.
     if "use_8bits_moe" in extra_options:
@@ -395,7 +397,10 @@ def set_io_dtype(precision, execution_provider, extra_options) -> ir.DataType:
     """
     Set the input/output precision of the ONNX model based on the provided precision and execution provider.
     """
-    cpu_quant = precision in {"int4", "int8"} and execution_provider == "cpu"
+    # int2/int4/int8 weight-only quantization builds a float graph and quantizes the weights at save time.
+    # On the CPU EP the I/O stays FP32; on GPU/WebGPU it follows the usual FP16 default (int8 must not
+    # be forced to FP32 I/O everywhere).
+    cpu_quant = precision in {"int2", "int4", "int8"} and execution_provider == "cpu"
     fp32_webgpu = execution_provider == "webgpu" and extra_options.get("use_webgpu_fp32", False)
     bf16_cuda = precision == "int4" and execution_provider in {"cuda", "trt-rtx"} and extra_options.get("use_cuda_bf16", False)
 
@@ -418,6 +423,11 @@ def set_onnx_dtype(precision: str, extra_options: dict[str, Any]) -> ir.DataType
     if precision == "int4":
         return ir.DataType.INT4 if extra_options.get("is_symmetric", True) else ir.DataType.UINT4
 
+    if precision == "int2":
+        # 2-bit quantized weights are emitted as MatMulNBits (per-module bits=2); non-quantized
+        # ops follow io_dtype, so onnx_dtype tracks the FLOAT io_dtype used on CPU.
+        return ir.DataType.FLOAT
+
     if precision == "int8":
         return ir.DataType.INT8 if extra_options.get("is_symmetric", True) else ir.DataType.UINT8
 
@@ -459,6 +469,10 @@ def create_model(
     # Set input/output precision of ONNX model
     io_dtype = set_io_dtype(precision, execution_provider, extra_options)
     onnx_dtype = set_onnx_dtype(precision, extra_options)
+    # int2 carries onnx_dtype FLOAT (2-bit weights are per-module MatMulNBits bits=2), so the
+    # onnx_dtype alone can't distinguish an int2 build from a genuine fp32 build. Thread the true
+    # precision string through so the builder can resolve the QMoE bit-width / op selection.
+    extra_options["precision"] = precision
     config_only = extra_options.get("config_only", False)
 
     # List architecture options in alphabetical order
@@ -484,6 +498,24 @@ def create_model(
         onnx_model = Gemma3Model(config, io_dtype, onnx_dtype, execution_provider, cache_dir, extra_options)
         if not onnx_model.exclude_embeds:
             onnx_model.model_type = "gemma3_vl_text"
+    elif config.architectures[0] == "Gemma4ForConditionalGeneration":
+        text_config = config.text_config
+        for key in text_config:
+            if not hasattr(config, key):
+                setattr(config, key, getattr(text_config, key))
+        print("WARNING: This model loses accuracy with float16 precision. It is recommended to set `--precision bf16` or `--precision int4 --extra_options use_cuda_bf16=true` by default.")
+        print("WARNING: This is only generating the text component of the model. The vision and audio components are not supported.")
+        onnx_model = Gemma4MoEModel(config, io_dtype, onnx_dtype, execution_provider, cache_dir, extra_options)
+        onnx_model.model_type = "gemma4_text"
+    elif config.architectures[0] == "Gemma4UnifiedForConditionalGeneration":
+        text_config = config.text_config
+        for key in text_config:
+            if not hasattr(config, key):
+                setattr(config, key, getattr(text_config, key))
+        print("WARNING: This model loses accuracy with float16 precision. It is recommended to set `--precision bf16` or `--precision int4 --extra_options use_cuda_bf16=true` by default.")
+        print("WARNING: This is only generating the text component of the model. The vision and audio components are not supported.")
+        onnx_model = Gemma4Model(config, io_dtype, onnx_dtype, execution_provider, cache_dir, extra_options)
+        onnx_model.model_type = "gemma4_text"
     elif config.architectures[0] == "GptOssForCausalLM":
         print("WARNING: This model only supports symmetric quantization for `QMoE`.")
         if hasattr(config, "quantization_config") and config.quantization_config.get("quant_method") != "quark":
@@ -620,7 +652,7 @@ def get_args():
         "-p",
         "--precision",
         required=True,
-        choices=["int4", "int8", "bf16", "fp16", "fp32"],
+        choices=["int2", "int4", "int8", "bf16", "fp16", "fp32"],
         help="Precision of model",
     )
 

src/python/py/models/builders/__init__.py

@@ -9,7 +9,7 @@
 from .base import Model
 from .chatglm import ChatGLMModel
 from .ernie import ErnieModel
-from .gemma import Gemma2Model, Gemma3Model, GemmaModel
+from .gemma import Gemma2Model, Gemma3Model, Gemma4MoEModel, Gemma4Model, GemmaModel
 from .gptoss import GPTOSSModel
 from .granite import GraniteModel, GraniteMoEHybridModel
 from .hunyuan import HunyuanDenseV1Model
@@ -52,6 +52,8 @@
     "GPTOSSModel",
     "Gemma2Model",
     "Gemma3Model",
+    "Gemma4MoEModel",
+    "Gemma4Model",
     "GemmaModel",
     "GraniteMoEHybridModel",
     "GraniteModel",

src/python/py/models/builders/base.py

@@ -89,11 +89,11 @@
             if hasattr(config, "num_hidden_layers")
             else config.num_layers
         )
         self.layer_types = (
             list(config.layer_types[: self.num_layers])
             if getattr(config, "layer_types", None) is not None
             else ["full_attention"] * self.num_layers
         )
         self.vocab_size = config.vocab_size
         self.activation = (
             config.hidden_activation
@@ -927,9 +927,13 @@
     def make_quant_config_init(self):
         self.quant_config = self.extra_options.get("_quant_config", None)
         if self.quant_config is None:
+            # int2 carries onnx_dtype FLOAT, so passing self.onnx_dtype would collapse it to "fp32"
+            # and lose the 2-bit signal. Prefer the true precision string threaded through
+            # extra_options (set by builder.create_model) when present.
+            quant_precision = self.extra_options.get("precision") or self.onnx_dtype
             self.quant_config = QuantConfig.from_extra_options(
                 extra_options=self.extra_options,
-                precision=self.onnx_dtype,
+                precision=quant_precision,
                 execution_provider=self.ep,
             )
         elif not isinstance(self.quant_config, QuantConfig):
@@ -951,7 +955,12 @@
 
         # MXFP4 and NVFP4 both resolve to the "mx" kind; the QMoE op tells them apart by dtype name
         # ("mxfp4" -> op "fp4", "nvfp4" -> op "nvfp4"). Integer dtypes use the plain "int" QMoE path.
-        self.moe_attrs["moe_op_type"] = "QMoE" if moe_descriptor.is_quantized else "MoE"
+        # Both keys carry the resolved op: the shared emitters (make_moe_expert_initializers /
+        # make_moe_op) read "op_type", while the gemma parallel-FFN path reads "moe_op_type". Setting
+        # only one left quantized experts on the float "MoE" op, which rejects uint8 weights.
+        moe_op_type = "QMoE" if moe_descriptor.is_quantized else "MoE"
+        self.moe_attrs["op_type"] = moe_op_type
+        self.moe_attrs["moe_op_type"] = moe_op_type
         if moe_descriptor.kind == "mx":
             self.moe_attrs["qmoe_quant_type"] = "nvfp4" if moe_descriptor.name == "nvfp4" else "fp4"
         else:
@@ -2057,6 +2066,30 @@
         )
 
     def make_matmul_op(self, matmul, basename, root_input, **kwargs):
+        # Factored online rotation (Quark rotation algo): rotate the activation before the
+        # quantized projection. Weights are stored in the rotated basis, so this is required.
+        # LoRA (below) is applied to the ORIGINAL (pre-rotation) activation, so capture it first.
+        original_root_input = root_input
+        if getattr(matmul, "input_prescale", None) is not None and getattr(self, "shared_input_rotations", None):
+            root_input = self.make_factored_rotation(matmul, basename, root_input, **kwargs)
+
+        matmul_name = self.make_matmul_core(matmul, basename, root_input, **kwargs)
+
+        # Bake the additive PEFT LoRA delta into the graph, if present (never for `logits`,
+        # since the excluded lm_head carries no adapter).
+        if (
+            not kwargs.get("logits", False)
+            and getattr(matmul, "lora_A", None) is not None
+            and getattr(matmul, "lora_B", None) is not None
+        ):
+            return self.make_lora_add(matmul, basename, original_root_input, matmul_name, **kwargs)
+        return matmul_name
+
+    def make_matmul_core(self, matmul, basename, root_input, **kwargs):
+        # 2-bit Quark weights are emitted as MatMulNBits even when the io/onnx dtype is
+        # float (ORT runs the float zero-point 2-bit MatMulNBits on CPU/MLAS).
+        if hasattr(matmul, "qweight") and matmul.qweight is not None and getattr(matmul, "bits", None) == 2:
+            return self.make_matmul_nbits(matmul, basename, root_input, **kwargs)
         if self.onnx_dtype in {ir.DataType.FLOAT16, ir.DataType.BFLOAT16, ir.DataType.FLOAT}:
             return self.make_matmul_float(matmul, basename, root_input, **kwargs)
         elif self.onnx_dtype in {ir.DataType.INT4, ir.DataType.UINT4, ir.DataType.INT8, ir.DataType.UINT8}:
@@ -2172,6 +2205,62 @@
         )
         return basename
 
+    def make_factored_rotation(self, matmul, basename, root_input, **kwargs):
+        """Emit the factored online input rotation for a projection:
+
+            x_rot = (x * input_prescale) @ shared_input_rotation_<in>
+
+        `input_prescale` is a per-projection [in] vector (small, emitted per projection).
+        `shared_input_rotation_<in>` is a single [in, in] matrix shared by every projection
+        of the same in_features and emitted only once (a huge size saving vs. inlining a copy
+        per projection). Returns the rotated activation tensor name to feed into MatMulNBits.
+        """
+        seq_dim = kwargs.get("seq_dim", "sequence_length")
+        in_features = matmul.in_features
+
+        # 1. Multiply by the per-projection input pre-scale vector.
+        prescale_init = basename[1:].replace("/", ".") + ".input_prescale"
+        self.make_initializer(matmul.input_prescale, prescale_init, to=self.io_dtype)
+        div_output = f"{basename}/input_prescale/output_0"
+        self.make_node("Mul", inputs=[root_input, prescale_init], outputs=[div_output], name=f"{basename}/input_prescale")
+        self.make_value(div_output, self.io_dtype, shape=["batch_size", seq_dim, in_features])
+
+        # 2. Multiply by the shared rotation matrix (single initializer per in_features).
+        rot_init = f"model.shared_input_rotation_{in_features}"
+        if rot_init not in self.shared_rotation_initializers:
+            self.make_initializer(self.shared_input_rotations[in_features], rot_init, to=self.io_dtype)
+            self.shared_rotation_initializers.add(rot_init)
+        rot_output = f"{basename}/shared_input_rotation/output_0"
+        self.make_node("MatMul", inputs=[div_output, rot_init], outputs=[rot_output], name=f"{basename}/shared_input_rotation")
+        self.make_value(rot_output, self.io_dtype, shape=["batch_size", seq_dim, in_features])
+        return rot_output
+
+    def make_lora_add(self, matmul, basename, root_input, matmul_name, **kwargs):
+        # Additive PEFT LoRA delta, baked into the graph:
+        #     y = MatMulNBits(rotate(x)) + ((x @ lora_A^T) @ (scaling * lora_B)^T)
+        # lora_A is [rank, in_features], lora_B is [out_features, rank]; both act on the
+        # ORIGINAL (pre-rotation) activation, matching how the adapter was trained. The
+        # `scaling` factor (lora_alpha / r) is pre-baked into the lora_B initializer.
+        seq_dim = kwargs.get("seq_dim", "sequence_length")
+        rank = matmul.lora_A.shape[0]
+        out_features = matmul.lora_B.shape[0]
+
+        lora_A_weight = f"{basename}/lora_A/weight"
+        self.make_initializer(matmul.lora_A.T, lora_A_weight, to=self.io_dtype)
+        lora_A_output = f"{basename}/lora_A/output_0"
+        self.make_node("MatMul", inputs=[root_input, lora_A_weight], outputs=[lora_A_output], name=f"{basename}/lora_A")
+        self.make_value(lora_A_output, self.io_dtype, shape=["batch_size", seq_dim, rank])
+
+        lora_B_weight = f"{basename}/lora_B/weight"
+        self.make_initializer((matmul.lora_B * matmul.lora_scaling).T, lora_B_weight, to=self.io_dtype)
+        lora_B_output = f"{basename}/lora_B/output_0"
+        self.make_node("MatMul", inputs=[lora_A_output, lora_B_weight], outputs=[lora_B_output], name=f"{basename}/lora_B")
+        self.make_value(lora_B_output, self.io_dtype, shape=["batch_size", seq_dim, out_features])
+
+        add_name = f"{basename}/lora/Add"
+        self.make_add(add_name, [f"{matmul_name}/output_0", lora_B_output], dtype=self.io_dtype, shape=["batch_size", seq_dim, out_features])
+        return add_name
+
     def make_matmul_float(self, matmul, name, root_input, **kwargs):
         weight = name[1:].replace("/", ".") + ".weight"
         self.make_initializer(matmul.weight.T, weight, to=self.io_dtype)
@@ -5175,6 +5264,13 @@
                 intermediate_size=self.intermediate_size,
                 num_layers=self.num_layers,
             )
+            # Factored online rotation (Quark rotation algo): shared [in, in] rotation
+            # matrices, emitted once and referenced by every projection of that in_features.
+            self.shared_input_rotations = getattr(model, "shared_input_rotations", {}) or {}
+            self.shared_rotation_initializers = set()
+            if self.shared_input_rotations:
+                # Each projection applies its own input rotation, so build q/k/v separately.
+                self.attention_attrs["use_packed_matmul"] = False
 
         else:
             extra_kwargs = {"num_hidden_layers": self.num_layers} if "num_hidden_layers" in self.extra_options else {}
@@ -5210,7 +5306,10 @@
                 **extra_kwargs,
             )
 
-        if "adapter_path" in self.extra_options:
+        # The quantized (Quark) loader attaches its LoRA adapter internally from
+        # <input_path>/lora_adapters/, and QuantModel is not an nn.Module, so the
+        # PeftModel wrapping only applies to the plain HF path.
+        if "adapter_path" in self.extra_options and self.quant_type is None:
             from peft import PeftModel
 
             model = PeftModel.from_pretrained(