From cd5fc2c7b92b6f88eba6047c4ee8fbb0b62ec8d0 Mon Sep 17 00:00:00 2001
From: Kunal Vaishnavi <kvaishnavi@microsoft.com>
Date: Tue, 15 Apr 2025 05:21:42 +0000
Subject: [PATCH 01/26] Initial implementation of bfloat16 support

---
 src/python/py/models/builder.py         | 277 ++++++++++++++----------
 src/python/py/models/quantized_model.py |   3 -
 2 files changed, 163 insertions(+), 117 deletions(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index 14e7e445dc..8e4cbcc6bd 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -15,6 +15,7 @@
 import torch
 
 import argparse
+import ctypes
 import gc
 import json
 import os
@@ -37,8 +38,8 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
 
         self.model_name_or_path = config._name_or_path
         self.model_type = config.architectures[0]
-        self.io_dtype = io_dtype      # {'fp16', 'fp32'}
-        self.onnx_dtype = onnx_dtype  # {"int4", "fp16", "fp32"}
+        self.io_dtype = io_dtype      # {"bf16", "fp16", "fp32"}
+        self.onnx_dtype = onnx_dtype  # {"int4", "bf16", "fp16", "fp32"}
         self.quant_type = config.quantization_config["quant_method"] if hasattr(config, "quantization_config") else None
         self.adapter_path = extra_options.get("adapter_path", None)
 
@@ -58,7 +59,7 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
         self.ep_attrs = {
             "cpu": {},
             "cuda": {
-                "enable_cuda_graph": "1" if extra_options.get("enable_cuda_graph", False) else "0",        # "1" if the model is able to enable cuda graph, "0" otherwise
+                "enable_cuda_graph": "1" if extra_options.get("enable_cuda_graph", False) else "0",              # "1" if the model is able to enable cuda graph, "0" otherwise
             },
             "rocm": {
                 "tunable_op_enable": "1",
@@ -123,11 +124,22 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
             TensorProto.FLOAT: np.float32,
         }
 
+        # Map TensorProto dtypes to PyTorch dtypes
+        self.to_torch_dtype = {
+            TensorProto.INT8: torch.int8,
+            TensorProto.INT32: torch.int32,
+            TensorProto.INT64: torch.int64,
+            TensorProto.BFLOAT16: torch.bfloat16,
+            TensorProto.FLOAT16: torch.float16,
+            TensorProto.FLOAT: torch.float32,
+        }
+
         # Map TensorProto dtypes to string dtypes
         self.to_str_dtype = {
             TensorProto.INT8: "TensorProto.INT8",
             TensorProto.INT32: "TensorProto.INT32",
             TensorProto.INT64: "TensorProto.INT64",
+            TensorProto.BFLOAT16: "TensorProto.BFLOAT16",
             TensorProto.FLOAT16: "TensorProto.FLOAT16",
             TensorProto.FLOAT: "TensorProto.FLOAT",
         }
@@ -280,12 +292,12 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
         # Quantization-specific variables (INT4, INT8, etc.)
         self.quant_attrs = {
             "int4": {
-                "accuracy_level": int(extra_options.get("int4_accuracy_level", 4 if self.ep == "cpu" else 0)),   # Default is 0 for non-QDQ formats, default is 4 for QDQ formats
+                "accuracy_level": int(extra_options.get("int4_accuracy_level", 4 if self.ep == "cpu" else 0)),
                 "block_size": int(extra_options.get("int4_block_size", 32)),
                 "is_symmetric": extra_options.get("int4_is_symmetric", True),
                 "op_types_to_quantize": extra_options.get("int4_op_types_to_quantize", ("MatMul", )),
             },
-            "use_qdq": extra_options.get("use_qdq", False),           # Use QDQ format
+            "use_qdq": extra_options.get("use_qdq", False),
         }
         if self.quant_type is not None:
             # Create quantized attributes from quantization config
@@ -295,7 +307,9 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
     def make_attention_init(self):
         valid_gqa_configurations = [
             ("cpu", TensorProto.FLOAT),
+            ("cpu", TensorProto.BFLOAT16),
             ("cuda", TensorProto.FLOAT16),
+            ("cuda", TensorProto.BFLOAT16),
             ("rocm", TensorProto.FLOAT16),
             ("dml", TensorProto.FLOAT16),
             ("webgpu", TensorProto.FLOAT16),
@@ -509,22 +523,33 @@ def order_repeated_field(self, repeated_proto, key_name, order):
         order = list(order)
         repeated_proto.sort(key=lambda x: order.index(getattr(x, key_name)))
 
-    def make_external_tensor(self, np_data, name, unpack_int4=False, **kwargs):
-        tensor = numpy_helper.from_array(np_data)
-        tensor.name = name
+    def make_external_tensor(self, tensor, name, **kwargs):
+        raw_data = bytes(
+            (ctypes.c_ubyte * tensor.element_size() * tensor.numel()).from_address(
+                tensor.data_ptr()
+            )
+        )
+        tensor_proto = helper.make_tensor(
+            name=name,
+            data_type=torch.onnx._type_utils.JitScalarType.from_dtype(tensor.dtype).onnx_type(),
+            dims=tensor.shape,
+            vals=raw_data,
+            raw=True,
+        )
 
         filename = f"{name}.bin"
-        external_data_helper.set_external_data(tensor, location=filename)
+        external_data_helper.set_external_data(tensor_proto, location=filename)
         with open(os.path.join(self.cache_dir, filename), "wb") as f:
-            f.write(tensor.raw_data)
-        tensor.ClearField("raw_data")
-        tensor.data_location = TensorProto.EXTERNAL
+            f.write(tensor_proto.raw_data)
+        tensor_proto.ClearField("raw_data")
+        tensor_proto.data_location = TensorProto.EXTERNAL
 
-        if unpack_int4 and self.onnx_dtype == 'int4':
-            tensor.data_type = TensorProto.UINT4
-            tensor.dims[-1] *= 2
+        # TODO: is this still needed for DML?
+        if kwargs.get("unpack_int4", False) and self.onnx_dtype == 'int4':
+            tensor_proto.data_type = TensorProto.UINT4
+            tensor_proto.dims[-1] *= 2
 
-        self.initializers.append(tensor)
+        self.initializers.append(tensor_proto)
 
     def make_node(self, op_type, inputs, outputs, name=None, doc_string=None, domain=None, **kwargs):
         # Save any constants as nodes
@@ -598,8 +623,21 @@ def make_constant(self, name):
         # Format of name is "/model/constants/{dtype}/{shape}/{num}"
         path = name.split("/")
         onnx_dtype, dims, num = eval(path[-3]), path[-2], eval(path[-1])
-        np_dtype = self.to_numpy_dtype[onnx_dtype]
-        value = numpy_helper.from_array(np.array(num if dims == "0D" else list(num) if type(num) == tuple else [num], dtype=np_dtype), name=name.replace("constants", "numpy_helper"))
+
+        torch_dtype = self.to_torch_dtype[onnx_dtype]
+        tensor = torch.tensor(num if dims == "0D" else list(num) if type(num) == tuple else [num], dtype=torch_dtype).contiguous()
+        raw_data = bytes(
+            (ctypes.c_ubyte * tensor.element_size() * tensor.numel()).from_address(
+                tensor.data_ptr()
+            )
+        )
+        value = helper.make_tensor(
+            name=name.replace("constants", "torch_helper"),
+            data_type=torch.onnx._type_utils.JitScalarType.from_dtype(tensor.dtype).onnx_type(),
+            dims=tensor.shape,
+            vals=raw_data,
+            raw=True,
+        )
 
         node_name = name.replace("constants", "constant_nodes")
         self.make_node("Constant", inputs=[], outputs=[name], name=node_name, value=value)
@@ -750,8 +788,8 @@ def make_matmul(self, matmul, basename, root_input, **kwargs):
             return self.make_matmul_op(matmul, basename, root_input, **kwargs)
     
     def make_matmul_op(self, matmul, basename, root_input, **kwargs):
-        if self.onnx_dtype in {"fp16", "fp32"}:
-            return self.make_matmul_fp16_or_fp32(matmul, basename, root_input, **kwargs)
+        if self.onnx_dtype in {"bf16", "fp16", "fp32"}:
+            return self.make_matmul_float(matmul, basename, root_input, **kwargs)
         elif self.onnx_dtype == "int4":
             if self.quant_attrs["use_qdq"]:
                 return self.make_matmul_int4_qdq(matmul, basename, root_input, **kwargs)
@@ -760,9 +798,9 @@ def make_matmul_op(self, matmul, basename, root_input, **kwargs):
         else:
             raise NotImplementedError(f"The {self.onnx_dtype} precision is not currently supported.")
 
-    def make_matmul_fp16_or_fp32(self, matmul, name, root_input, **kwargs):
+    def make_matmul_float(self, matmul, name, root_input, **kwargs):
         weight = name[1:].replace("/", ".") + ".weight"
-        self.make_external_tensor(matmul.weight.detach().numpy().transpose().astype(self.to_numpy_dtype[self.io_dtype]), weight)
+        self.make_external_tensor(matmul.weight.detach().T.to(self.to_torch_dtype[self.io_dtype]).contiguous(), weight)
 
         last_dim = matmul.weight.shape[0]
         output = "logits" if kwargs.get("logits", False) else f"{name}/output_0"
@@ -773,29 +811,29 @@ def make_matmul_fp16_or_fp32(self, matmul, name, root_input, **kwargs):
 
     def make_matmul_int4(self, matmul, basename, root_input, **kwargs):
         if not hasattr(matmul, "qweight"):
-            # TODO: quantize weights, then save new MatMul numpy weights for onnx model
+            # TODO: quantize weights, then save new MatMul weights for onnx model
             # print(f"Quantizing to {self.onnx_dtype} on-the-fly is not currently supported.")
             # print(f"Saving as {self.io_dtype} on-the-fly and quantizing to {self.onnx_dtype} at the end.")
-            return self.make_matmul_fp16_or_fp32(matmul, basename, root_input, **kwargs)
+            return self.make_matmul_float(matmul, basename, root_input, **kwargs)
 
         name = f"{basename}NBits"
 
-        # Input weights are quantized, save quantized MatMul numpy weights for onnx model
+        # Input weights are quantized, save quantized MatMul weights for onnx model
         weight = name[1:].replace("/", ".") + ".qweight"
-        self.make_external_tensor(matmul.qweight.detach().numpy(), weight)
+        self.make_external_tensor(matmul.qweight.detach().contiguous(), weight)
         scales = name[1:].replace("/", ".") + ".scales"
-        self.make_external_tensor(matmul.scales.detach().numpy().astype(self.to_numpy_dtype[self.io_dtype]), scales)
+        self.make_external_tensor(matmul.scales.detach().to(self.to_torch_dtype[self.io_dtype]).contiguous(), scales)
 
         inputs = [root_input, weight, scales]
 
         if hasattr(matmul, "qzeros") and matmul.qzeros is not None:
             zeros = name[1:].replace("/", ".") + ".qzeros"
-            self.make_external_tensor(matmul.qzeros.detach().numpy(), zeros)
+            self.make_external_tensor(matmul.qzeros.detach().contiguous(), zeros)
             inputs.append(zeros)
 
         if hasattr(matmul, "g_idx") and matmul.g_idx is not None:
             g_idx = name[1:].replace("/", ".") + ".g_idx"
-            self.make_external_tensor(matmul.g_idx.detach().numpy().astype(np.int32), g_idx)
+            self.make_external_tensor(matmul.g_idx.detach().to(torch.int32).contiguous(), g_idx)
             inputs.append(g_idx)
 
         output = "logits" if kwargs.get("logits", False) else f"{name}/output_0"
@@ -808,25 +846,26 @@ def make_matmul_int4(self, matmul, basename, root_input, **kwargs):
 
         return name
 
+    # TODO: are the booleans in make_external_tensor still needed?
     def make_dequantize_linear(self, dequantize_name, quantized_op):
-        # Input weights are quantized, save quantized MatMul numpy weights for onnx model
+        # Input weights are quantized, save quantized MatMul weights for onnx model
         qweight = dequantize_name[1:].replace("/", ".") + ".qweight"
-        qweight_npy = quantized_op.qweight.detach().numpy()
+        qweight_npy = quantized_op.qweight.detach()
         qweight_npy = qweight_npy.reshape(*qweight_npy.shape[:-2], qweight_npy.shape[-2] * qweight_npy.shape[-1])
-        self.make_external_tensor(qweight_npy, qweight, True)
+        self.make_external_tensor(qweight_npy.contiguous(), qweight, True)
 
         scales = dequantize_name[1:].replace("/", ".") + ".scales"
-        scales_npy = quantized_op.scales.detach().numpy().astype(self.to_numpy_dtype[self.io_dtype])
+        scales_npy = quantized_op.scales.detach().to(self.to_torch_dtype[self.io_dtype])
         scales_npy = scales_npy.reshape(*qweight_npy.shape[:-1], qweight_npy.shape[-1] * 2 // quantized_op.group_size)
-        self.make_external_tensor(scales_npy, scales)
+        self.make_external_tensor(scales_npy.contiguous(), scales)
 
         dequantize_inputs = [qweight, scales]
 
         if hasattr(quantized_op, "qzeros") and quantized_op.qzeros is not None:
             zeros = dequantize_name[1:].replace("/", ".") + ".qzeros"
-            zeros_npy = quantized_op.qzeros.detach().numpy()
+            zeros_npy = quantized_op.qzeros.detach()
             zeros_npy = zeros_npy.reshape(*qweight_npy.shape[:-1], qweight_npy.shape[-1] // quantized_op.group_size)
-            self.make_external_tensor(zeros_npy, zeros, True)
+            self.make_external_tensor(zeros_npy.contiguous(), zeros, True)
             dequantize_inputs.append(zeros)
 
         dequantize_output = f"{dequantize_name}/output_0"
@@ -837,10 +876,10 @@ def make_dequantize_linear(self, dequantize_name, quantized_op):
 
     def make_matmul_int4_qdq(self, matmul, matmul_name, root_input, **kwargs):
         if not hasattr(matmul, "qweight"):
-            # TODO: quantize weights, then save new MatMul numpy weights for onnx model
+            # TODO: quantize weights, then save new MatMul weights for onnx model
             # print(f"Quantizing to {self.onnx_dtype} on-the-fly is not currently supported.")
             # print(f"Saving as {self.io_dtype} on-the-fly and quantizing to {self.onnx_dtype} at the end.")
-            return self.make_matmul_fp16_or_fp32(matmul, matmul_name, root_input, **kwargs)
+            return self.make_matmul_float(matmul, matmul_name, root_input, **kwargs)
 
         dequantize_output = self.make_dequantize_linear(f"{matmul_name}/DequantizeLinear", matmul)
 
@@ -848,7 +887,7 @@ def make_matmul_int4_qdq(self, matmul, matmul_name, root_input, **kwargs):
         # compute quantized matmul when the weights are transposed. In most implementations, the transpose should usually be converted to a "transposeB"
         # attribute on the MatMul itself. A more natural way to represent this would have been to use Gemm since it already supports a transB attribute,
         # but unfortunately Gemm doesn't support batches.
-        qweight_shape = matmul.qweight.detach().numpy().shape
+        qweight_shape = matmul.qweight.detach().shape
         transposed_shape = [qweight_shape[1] * qweight_shape[2] * 2, qweight_shape[0]]
         transpose_name = f"{matmul_name}/Transpose"
         self.make_transpose(transpose_name, dequantize_output, self.io_dtype, transposed_shape, [1, 0])
@@ -900,14 +939,14 @@ def make_matmul_lora(self, matmul, basename, root_input, **kwargs):
         return add_name
 
     def make_packed_matmul(self, q_matmul, k_matmul, v_matmul, basename, root_input, **kwargs):
-        if self.onnx_dtype in {"fp16", "fp32"}:
-            return self.make_packed_matmul_fp16_or_fp32(q_matmul, k_matmul, v_matmul, basename, root_input, **kwargs)
+        if self.onnx_dtype in {"bf16", "fp16", "fp32"}:
+            return self.make_packed_matmul_float(q_matmul, k_matmul, v_matmul, basename, root_input, **kwargs)
         elif self.onnx_dtype == "int4":
             return self.make_packed_matmul_int4(q_matmul, k_matmul, v_matmul, basename, root_input, **kwargs)
         else:
             raise NotImplementedError(f"The {self.onnx_dtype} precision is not currently supported.")
 
-    def make_packed_matmul_fp16_or_fp32(self, q_matmul, k_matmul, v_matmul, name, root_input, **kwargs):
+    def make_packed_matmul_float(self, q_matmul, k_matmul, v_matmul, name, root_input, **kwargs):
         # N_q = num_attention_heads * head_size, N_kv = num_key_value_heads * head_size, H = hidden_size
         # Combine 3 MatMuls of shape N_q x H, N_kv x H, N_kv x H into 1 packed MatMul of shape (N_q+N_kv+N_kv)xH
         # Note: Packed MatMul is of shape (N_q+N_kv+N_kv)xH instead of Hx(N_q+N_kv+N_kv) because `make_matmul` will
@@ -918,7 +957,7 @@ def make_packed_matmul_fp16_or_fp32(self, q_matmul, k_matmul, v_matmul, name, ro
         # Create dummy PackedMatMul class
         class PackedMatMul:
             def __init__(self):
-                self.weight = torch.concatenate([q_matmul.weight.detach().cpu(), k_matmul.weight.detach().cpu(), v_matmul.weight.detach().cpu()], dim=0).reshape(N_q + N_kv + N_kv, H)
+                self.weight = torch.cat([q_matmul.weight.detach().cpu(), k_matmul.weight.detach().cpu(), v_matmul.weight.detach().cpu()], dim=0).reshape(N_q + N_kv + N_kv, H)
         matmul = PackedMatMul()
         new_name = self.make_matmul(matmul, name, root_input, **kwargs)
 
@@ -926,19 +965,19 @@ def __init__(self):
 
     def make_packed_matmul_int4(self, q_matmul, k_matmul, v_matmul, basename, root_input, **kwargs):
         if not hasattr(q_matmul, "qweight"):
-            # TODO: quantize weights, then save new MatMul numpy weights for onnx model
+            # TODO: quantize weights, then save new MatMul weights for onnx model
             # print(f"Quantizing to {self.onnx_dtype} on-the-fly is not currently supported.")
             # print(f"Saving as {self.io_dtype} on-the-fly and quantizing to {self.onnx_dtype} at the end.")
-            return self.make_packed_matmul_fp16_or_fp32(q_matmul, k_matmul, v_matmul, basename, root_input, **kwargs)
+            return self.make_packed_matmul_float(q_matmul, k_matmul, v_matmul, basename, root_input, **kwargs)
 
         name = f"{basename}NBits"
 
         # Create dummy PackedMatMul class
         class PackedMatMul:
             def __init__(self):
-                self.qweight = torch.concatenate([q_matmul.qweight.detach().cpu(), k_matmul.qweight.detach().cpu(), v_matmul.qweight.detach().cpu()], dim=0)
-                self.scales = torch.concatenate([q_matmul.scales.detach().cpu(), k_matmul.scales.detach().cpu(), v_matmul.scales.detach().cpu()], dim=0)
-                self.qzeros = torch.concatenate([q_matmul.qzeros.detach().cpu(), k_matmul.qzeros.detach().cpu(), v_matmul.qzeros.detach().cpu()], dim=0)
+                self.qweight = torch.cat([q_matmul.qweight.detach().cpu(), k_matmul.qweight.detach().cpu(), v_matmul.qweight.detach().cpu()], dim=0)
+                self.scales = torch.cat([q_matmul.scales.detach().cpu(), k_matmul.scales.detach().cpu(), v_matmul.scales.detach().cpu()], dim=0)
+                self.qzeros = torch.cat([q_matmul.qzeros.detach().cpu(), k_matmul.qzeros.detach().cpu(), v_matmul.qzeros.detach().cpu()], dim=0)
                 self.g_idx = q_matmul.g_idx
 
                 self.in_features = q_matmul.in_features
@@ -947,22 +986,22 @@ def __init__(self):
                 self.group_size = q_matmul.group_size
         matmul = PackedMatMul()
 
-        # Input weights are quantized, save quantized MatMul numpy weights for onnx model
+        # Input weights are quantized, save quantized MatMul weights for onnx model
         weight = name[1:].replace("/", ".") + ".qweight"
-        self.make_external_tensor(matmul.qweight.detach().numpy(), weight)
+        self.make_external_tensor(matmul.qweight.detach().contiguous(), weight)
         scales = name[1:].replace("/", ".") + ".scales"
-        self.make_external_tensor(matmul.scales.detach().numpy().astype(self.to_numpy_dtype[self.io_dtype]), scales)
+        self.make_external_tensor(matmul.scales.detach().to(self.to_torch_dtype[self.io_dtype]).contiguous(), scales)
 
         inputs = [root_input, weight, scales]
 
         if hasattr(matmul, "qzeros") and matmul.qzeros is not None:
             zeros = name[1:].replace("/", ".") + ".qzeros"
-            self.make_external_tensor(matmul.qzeros.detach().numpy(), zeros)
+            self.make_external_tensor(matmul.qzeros.detach().contiguous(), zeros)
             inputs.append(zeros)
 
         if hasattr(matmul, "g_idx") and matmul.g_idx is not None:
             g_idx = name[1:].replace("/", ".") + ".g_idx"
-            self.make_external_tensor(matmul.g_idx.detach().numpy().astype(np.int32), g_idx)
+            self.make_external_tensor(matmul.g_idx.detach().to(torch.int32).contiguous(), g_idx)
             inputs.append(g_idx)
 
         output = "logits" if kwargs.get("logits", False) else f"{name}/output_0"
@@ -977,7 +1016,7 @@ def __init__(self):
 
     def make_add_bias(self, add, name, root_input, **kwargs):
         bias = name[1:].replace("/", ".") + ".bias"
-        self.make_external_tensor(add.astype(self.to_numpy_dtype[self.io_dtype]), bias)
+        self.make_external_tensor(add.to(self.to_torch_dtype[self.io_dtype]).contiguous(), bias)
 
         add_bias_inputs = [root_input, bias]
         shape = ['batch_size', 'sequence_length', add.shape[0]]
@@ -991,12 +1030,12 @@ def make_add_bias(self, add, name, root_input, **kwargs):
 
     def make_packed_add(self, q_add, k_add, v_add, name, root_input, **kwargs):
         # Combine 3 Adds of shape N_q, N_kv, and N_kv into 1 packed Add of shape N_q + N_kv + N_kv
-        add = np.concatenate([q_add, k_add, v_add], axis=0).flatten()
+        add = torch.cat([q_add, k_add, v_add], dim=0).flatten()
         self.make_add_bias(add, name, root_input, **kwargs)
 
     def make_embedding(self, embedding):
         weight = "model.embed_tokens.weight"
-        self.make_external_tensor(embedding.astype(self.to_numpy_dtype[self.io_dtype]), weight)
+        self.make_external_tensor(embedding.to(self.to_torch_dtype[self.io_dtype]).contiguous(), weight)
 
         basename = "/model/embed_tokens"
         gather_name = f"{basename}/Gather"
@@ -1024,10 +1063,10 @@ def make_layernorm(self, layer_id, layernorm, skip, simple, location):
         skip_input = self.layernorm_attrs["skip_input"]
 
         weight = f"model.layers.{layer_id}.{location}_layernorm.weight"
-        self.make_external_tensor(layernorm.weight.detach().numpy().astype(self.to_numpy_dtype[self.io_dtype]) + self.layernorm_attrs["add_offset"], weight)
+        self.make_external_tensor(layernorm.weight.detach().to(self.to_torch_dtype[self.io_dtype]).contiguous() + self.layernorm_attrs["add_offset"], weight)
         bias = f"model.layers.{layer_id}.{location}_layernorm.bias"
         if not simple:
-            self.make_external_tensor(layernorm.bias.detach().numpy().astype(self.to_numpy_dtype[self.io_dtype]), bias)
+            self.make_external_tensor(layernorm.bias.detach().to(self.to_torch_dtype[self.io_dtype]).contiguous(), bias)
 
         inputs = [root_input, skip_input, weight] if skip else [root_input, weight]
         if not simple:
@@ -1121,10 +1160,10 @@ def make_rotary_embedding_caches(self, **kwargs):
 
             # Slice cos/sin caches from (M, H) to (M, H/2)
             hidden_dim = cos_cache.shape[-1]
-            cos_cache = cos_cache.squeeze()[:, : (hidden_dim // 2)].detach().numpy()
-            cos_cache = cos_cache.astype(self.to_numpy_dtype[self.io_dtype])
-            sin_cache = sin_cache.squeeze()[:, : (hidden_dim // 2)].detach().numpy()
-            sin_cache = sin_cache.astype(self.to_numpy_dtype[self.io_dtype])
+            cos_cache = cos_cache.squeeze()[:, : (hidden_dim // 2)].detach()
+            cos_cache = cos_cache.to(self.to_torch_dtype[self.io_dtype])
+            sin_cache = sin_cache.squeeze()[:, : (hidden_dim // 2)].detach()
+            sin_cache = sin_cache.to(self.to_torch_dtype[self.io_dtype])
 
             # Slice cos/sin caches from (M, H/2) to (M, R/2) if partial rotary embeddings are used
             if self.rotemb_attrs["partial_rotary_factor"] != 1.0:
@@ -1133,8 +1172,8 @@ def make_rotary_embedding_caches(self, **kwargs):
 
             if self.rotemb_attrs["save_caches"]:
                 # Save cos/sin caches to disk
-                self.make_external_tensor(cos_cache, cos_cache_name)
-                self.make_external_tensor(sin_cache, sin_cache_name)
+                self.make_external_tensor(cos_cache.contiguous(), cos_cache_name)
+                self.make_external_tensor(sin_cache.contiguous(), sin_cache_name)
             else:
                 # Return cos/sin caches since they will be custom-saved
                 return cos_cache, sin_cache
@@ -1214,6 +1253,7 @@ def make_rotary_embedding_multi_cache(self):
                 initializer=[],
                 value_info=[],
                 nodes=[
+                    # TODO: fix constant node values since caches can be bf16
                     helper.make_node("Constant", inputs=[], outputs=[cos_cache_large_name], name="/large/cos_cache/Constant", value=numpy_helper.from_array(cos_cache_large)),
                     helper.make_node("Constant", inputs=[], outputs=[sin_cache_large_name], name="/large/sin_cache/Constant", value=numpy_helper.from_array(sin_cache_large)),
                 ],
@@ -1228,6 +1268,7 @@ def make_rotary_embedding_multi_cache(self):
                 initializer=[],
                 value_info=[],
                 nodes=[
+                    # TODO: fix constant node values since caches can be bf16
                     helper.make_node("Constant", inputs=[], outputs=[cos_cache_small_name], name="/small/cos_cache/Constant", value=numpy_helper.from_array(cos_cache_small)),
                     helper.make_node("Constant", inputs=[], outputs=[sin_cache_small_name], name="/small/sin_cache/Constant", value=numpy_helper.from_array(sin_cache_small)),
                 ],
@@ -1260,7 +1301,7 @@ def make_qk_norm(self, layer_id, attention):
         q_layernorm_name = f"/model/layers.{layer_id}/attn/q_norm/SimplifiedLayerNormalization"
         q_weight_name = f"model.layers.{layer_id}.attn.q_norm.layernorm.weight"
         q_layernorm_output = f"{q_layernorm_name}/output_0"
-        self.make_external_tensor(attention.q_norm.weight.detach().numpy().astype(self.to_numpy_dtype[self.io_dtype]) + self.layernorm_attrs["add_offset"], q_weight_name)
+        self.make_external_tensor((attention.q_norm.weight.detach().to(self.to_torch_dtype[self.io_dtype]) + self.layernorm_attrs["add_offset"]).contiguous(), q_weight_name)
         self.make_node("SimplifiedLayerNormalization", inputs=[q_reshape_1_output, q_weight_name], outputs=[q_layernorm_output], name=q_layernorm_name, **layernorm_kwargs)
         self.make_value_info(q_layernorm_output, dtype=self.io_dtype, shape=['batch_size', 'sequence_length * num_attention_heads', self.head_size])
 
@@ -1279,7 +1320,7 @@ def make_qk_norm(self, layer_id, attention):
         k_layernorm_name = f"/model/layers.{layer_id}/attn/k_norm/SimplifiedLayerNormalization"
         k_weight_name = f"model.layers.{layer_id}.attn.k_norm.layernorm.weight"
         k_layernorm_output = f"{k_layernorm_name}/output_0"
-        self.make_external_tensor(attention.k_norm.weight.detach().numpy().astype(self.to_numpy_dtype[self.io_dtype]) + self.layernorm_attrs["add_offset"], k_weight_name)
+        self.make_external_tensor((attention.k_norm.weight.detach().to(self.to_torch_dtype[self.io_dtype]) + self.layernorm_attrs["add_offset"]).contiguous(), k_weight_name)
         self.make_node("SimplifiedLayerNormalization", inputs=[k_reshape_1_output, k_weight_name], outputs=[k_layernorm_output], name=k_layernorm_name, **layernorm_kwargs)
         self.make_value_info(k_layernorm_output, dtype=self.io_dtype, shape=['batch_size', 'sequence_length * num_key_value_heads', self.head_size])
 
@@ -1592,20 +1633,20 @@ def make_attention(self, layer_id, attention, root_input, **kwargs):
         if all_bias_exists and self.attention_attrs["use_packed_matmul"]:
             # Combine 3 Adds into 1 packed Add
             qkv_add_name = f"/model/layers.{layer_id}/attn/qkv_proj/Add"
-            self.make_packed_add(attention.q_proj.bias.detach().numpy(), attention.k_proj.bias.detach().numpy(), attention.v_proj.bias.detach().numpy(), qkv_add_name, root_input=self.attention_attrs["q_path"])
+            self.make_packed_add(attention.q_proj.bias.detach(), attention.k_proj.bias.detach(), attention.v_proj.bias.detach(), qkv_add_name, root_input=self.attention_attrs["q_path"])
             self.attention_attrs["q_path"] = f"{qkv_add_name}/output_0"
         else:
             if q_bias_exists:
                 q_add_name = f"/model/layers.{layer_id}/attn/q_proj/Add"
-                self.make_add_bias(attention.q_proj.bias.detach().numpy(), q_add_name, root_input=self.attention_attrs["q_path"])
+                self.make_add_bias(attention.q_proj.bias.detach(), q_add_name, root_input=self.attention_attrs["q_path"])
                 self.attention_attrs["q_path"] = f"{q_add_name}/output_0"
             if k_bias_exists:
                 k_add_name = f"/model/layers.{layer_id}/attn/k_proj/Add"
-                self.make_add_bias(attention.k_proj.bias.detach().numpy(), k_add_name, root_input=self.attention_attrs["k_path"])
+                self.make_add_bias(attention.k_proj.bias.detach(), k_add_name, root_input=self.attention_attrs["k_path"])
                 self.attention_attrs["k_path"] = f"{k_add_name}/output_0"
             if v_bias_exists:
                 v_add_name = f"/model/layers.{layer_id}/attn/v_proj/Add"
-                self.make_add_bias(attention.v_proj.bias.detach().numpy(), v_add_name, root_input=self.attention_attrs["v_path"])
+                self.make_add_bias(attention.v_proj.bias.detach(), v_add_name, root_input=self.attention_attrs["v_path"])
                 self.attention_attrs["v_path"] = f"{v_add_name}/output_0"
 
         # Make Q/K SimplifiedLayerNorm nodes
@@ -1652,7 +1693,7 @@ def make_attention(self, layer_id, attention, root_input, **kwargs):
         o_bias_exists = eval(f"attention.{o_proj}.bias") is not None
         if o_bias_exists:
             o_add_name = f"/model/layers.{layer_id}/attn/o_proj/Add"
-            o_bias = eval(f"attention.{o_proj}.bias.detach().numpy()")
+            o_bias = eval(f"attention.{o_proj}.bias.detach()")
             self.make_add_bias(o_bias, o_add_name, root_input=f"{o_matmul_name}/output_0")
 
         # Assign output 0 of previous output node as skip input to next SkipLayerNorm
@@ -1841,7 +1882,7 @@ def make_mlp_proj(self, layer_id, mlp, root_input):
         gate_name = gate_matmul_name
         if gate_bias_exists:
             gate_add_name = f"/model/layers.{layer_id}/mlp/gate_proj/Add"
-            self.make_add_bias(mlp.gate_proj.bias.detach().numpy(), gate_add_name, root_input=f"{gate_name}/output_0")
+            self.make_add_bias(mlp.gate_proj.bias.detach(), gate_add_name, root_input=f"{gate_name}/output_0")
             gate_name = gate_add_name
 
         # Make Up proj nodes
@@ -1850,7 +1891,7 @@ def make_mlp_proj(self, layer_id, mlp, root_input):
         up_name = up_matmul_name
         if up_bias_exists:
             up_add_name = f"/model/layers.{layer_id}/mlp/up_proj/Add"
-            self.make_add_bias(mlp.up_proj.bias.detach().numpy(), up_add_name, root_input=f"{up_name}/output_0")
+            self.make_add_bias(mlp.up_proj.bias.detach(), up_add_name, root_input=f"{up_name}/output_0")
             up_name = up_add_name
 
         # Make activation node(s)
@@ -1867,7 +1908,7 @@ def make_mlp_proj(self, layer_id, mlp, root_input):
         down_name = down_matmul_name
         if down_bias_exists:
             down_add_name = f"/model/layers.{layer_id}/mlp/down_proj/Add"
-            self.make_add_bias(mlp.down_proj.bias.detach().numpy(), down_add_name, root_input=f"{down_name}/output_0")
+            self.make_add_bias(mlp.down_proj.bias.detach(), down_add_name, root_input=f"{down_name}/output_0")
             down_name = down_add_name
 
         # Assign output 0 of previous MatMul as skip input to next SkipLayerNorm
@@ -1898,7 +1939,7 @@ def make_mlp_fc(self, layer_id, mlp, root_input):
         fc1_name = fc1_matmul_name
         if fc1_bias_exists:
             fc1_add_name = f"/model/layers.{layer_id}/mlp/fc1/Add"
-            self.make_add_bias(mlp.fc1.bias.detach().numpy(), fc1_add_name, root_input=f"{fc1_name}/output_0")
+            self.make_add_bias(mlp.fc1.bias.detach(), fc1_add_name, root_input=f"{fc1_name}/output_0")
             fc1_name = fc1_add_name
 
         # Make activation function
@@ -1910,7 +1951,7 @@ def make_mlp_fc(self, layer_id, mlp, root_input):
         fc2_name = fc2_matmul_name
         if fc2_bias_exists:
             fc2_add_name = f"/model/layers.{layer_id}/mlp/fc2/Add"
-            self.make_add_bias(mlp.fc2.bias.detach().numpy(), fc2_add_name, root_input=f"{fc2_name}/output_0")
+            self.make_add_bias(mlp.fc2.bias.detach(), fc2_add_name, root_input=f"{fc2_name}/output_0")
             fc2_name = fc2_add_name
 
         # Assign output 0 of MLP layer as output of last layer
@@ -1949,19 +1990,14 @@ def make_block_sparse_moe(self, layer_id, bsm, root_input):
         # Make MoE nodes
         gate_name = f"{gate_ops_base}/MatMul"
         self.make_matmul(bsm.gate, gate_name, root_input)
-
         shape_name = f"{gate_ops_base}/Shape"
         self.make_shape(shape_name, f"{gate_name}/output_0", shape=[3])
-
         gather_name = f"{gate_ops_base}/Gather"
         self.make_gather(gather_name, [f"{shape_name}/output_0", "/model/constants/TensorProto.INT64/0D/2"], axis=0)
-
         unsqueeze_name = f"{gate_ops_base}/Unsqueeze"
         self.make_unsqueeze(unsqueeze_name, [f"{gather_name}/output_0", "/model/constants/TensorProto.INT64/1D/0"], dtype=TensorProto.INT64, shape=[1])
-
         concat_name = f"{gate_ops_base}/Concat"
         self.make_concat(concat_name, ["/model/constants/TensorProto.INT64/1D/-1", f"{unsqueeze_name}/output_0"], dtype=TensorProto.INT64, shape=[2], axis=0)
-
         gate_reshape_name = f"{gate_ops_base}/Reshape"
         self.make_reshape(gate_reshape_name, [f"{gate_name}/output_0", f"{concat_name}/output_0"], dtype=self.io_dtype, shape=['num_rows', num_experts])
 
@@ -2009,18 +2045,18 @@ def quant_dequant(weights, quant_mode: bool = True):
         moe_expert_scales_2_name = f"model.layers.{layer_id}.moe.scales_2"
         moe_expert_scales_3_name = f"model.layers.{layer_id}.moe.scales_3"
 
-        def make_moe_external_tensor(w_list, moe_expert_name, numpy_type):
-            moe_experts_weight = torch.stack(w_list, dim=0).detach().numpy()
-            self.make_external_tensor(moe_experts_weight.astype(numpy_type), moe_expert_name)
+        def make_moe_external_tensor(w_list, moe_expert_name, dtype):
+            moe_experts_weight = torch.stack(w_list, dim=0).detach()
+            self.make_external_tensor(moe_experts_weight.to(dtype).contiguous(), moe_expert_name)
 
-        make_moe_external_tensor(w1_list, moe_expert_weight_1_name, np.uint8)
-        make_moe_external_tensor(w2_list, moe_expert_weight_2_name, np.uint8)
-        make_moe_external_tensor(w3_list, moe_expert_weight_3_name, np.uint8)
+        make_moe_external_tensor(w1_list, moe_expert_weight_1_name, torch.uint8)
+        make_moe_external_tensor(w2_list, moe_expert_weight_2_name, torch.uint8)
+        make_moe_external_tensor(w3_list, moe_expert_weight_3_name, torch.uint8)
 
         # Currently we don't expect QMoE to be used with distributed inference
-        make_moe_external_tensor(w1_scale_list, moe_expert_scales_1_name, self.to_numpy_dtype[self.io_dtype])
-        make_moe_external_tensor(w2_scale_list, moe_expert_scales_2_name, self.to_numpy_dtype[self.io_dtype])
-        make_moe_external_tensor(w3_scale_list, moe_expert_scales_3_name, self.to_numpy_dtype[self.io_dtype])
+        make_moe_external_tensor(w1_scale_list, moe_expert_scales_1_name, self.to_torch_dtype[self.io_dtype])
+        make_moe_external_tensor(w2_scale_list, moe_expert_scales_2_name, self.to_torch_dtype[self.io_dtype])
+        make_moe_external_tensor(w3_scale_list, moe_expert_scales_3_name, self.to_torch_dtype[self.io_dtype])
 
         bias_ph = "" # Placeholder for bias
         inputs = [root_input, f"{gate_reshape_name}/output_0", \
@@ -2033,7 +2069,6 @@ def make_moe_external_tensor(w_list, moe_expert_name, numpy_type):
         self.make_node(op_type, inputs=inputs, outputs=[output], name=moe_name, domain="com.microsoft",
                        k=top_k, activation_type=activation_type, normalize_routing_weights=normalize_routing_weights,
                        use_sparse_mixer=use_sparse_mixer, expert_weight_bits=(4 if use_int4 else 8))
-
         self.make_value_info(output, self.io_dtype, shape=['batch_size', 'sequence_length', self.hidden_size])
 
         # Assign output 0 of previous MoE as root input to next SkipLayerNorm
@@ -2118,7 +2153,7 @@ def make_lm_head(self, lm_head):
 
         if bias_exists:
             add_name = "/lm_head/Add"
-            self.make_add_bias(lm_head.bias.detach().numpy(), add_name, root_input=f"{lm_name}/output_0", logits=not(scale_exists or mask_exists or softcap_exists))
+            self.make_add_bias(lm_head.bias.detach(), add_name, root_input=f"{lm_name}/output_0", logits=not(scale_exists or mask_exists or softcap_exists))
             lm_name = add_name
 
         if scale_exists:
@@ -2132,10 +2167,10 @@ def make_lm_head(self, lm_head):
         if mask_exists:
             # Save logits mask as initializer
             logits_mask_name = "logits_mask"
-            self.make_external_tensor(self.lm_head_attrs["mask"].detach().numpy(), logits_mask_name)
+            self.make_external_tensor(self.lm_head_attrs["mask"].detach().contiguous(), logits_mask_name)
 
             where_name = "/lm_head/Where"
-            where_inputs = [logits_mask_name, f"/model/constants/{self.to_str_dtype[self.io_dtype]}/0D/{np.finfo(self.to_numpy_dtype[self.io_dtype]).min}", f"{lm_name}/output_0"]
+            where_inputs = [logits_mask_name, f"/model/constants/{self.to_str_dtype[self.io_dtype]}/0D/{torch.finfo(self.to_torch_dtype[self.io_dtype]).min}", f"{lm_name}/output_0"]
             where_output = "logits" if not softcap_exists else f"{where_name}/output_0"
             self.make_node('Where', inputs=where_inputs, outputs=[where_output], name=where_name)
             self.make_value_info(where_output, self.io_dtype, shape=['batch_size', 'sequence_length', self.vocab_size])
@@ -2220,7 +2255,7 @@ def make_model(self, input_path):
                 if not self.exclude_embeds:
                     # Embedding layer
                     print("Reading embedding layer")
-                    self.make_embedding(module.weight.detach().numpy())
+                    self.make_embedding(module.weight.detach())
                 else:
                     # Exclude embedding layer from model
                     self.layernorm_attrs["root_input"] = "inputs_embeds"
@@ -2410,6 +2445,7 @@ def make_input_ids_subgraph(self, basename, past_key_gather_name):
         concat_2_name = f"{basename}/Concat_2"
         concat_inputs = [f"{unsqueeze_4_name}/output_0", f"{unsqueeze_5_name}/output_0"]
         self.make_concat(concat_2_name, concat_inputs, dtype=TensorProto.INT64, shape=[2], axis=0)
+        # TODO: use to_torch_dtype for constant node
         constant_shape_name = f"{basename}/ConstantOfShape_2"
         constant_shape_numpy_dtype = self.to_numpy_dtype[self.io_dtype]
         constant_shape_value = numpy_helper.from_array(np.array([np.finfo(constant_shape_numpy_dtype).min], dtype=constant_shape_numpy_dtype))
@@ -2419,7 +2455,7 @@ def make_input_ids_subgraph(self, basename, past_key_gather_name):
         shape_4_name = f"{basename}/Shape_4"
         self.make_shape(shape_4_name, f"{constant_shape_name}/output_0", shape=[2])
         slice_1_name = f"{basename}/Slice_1"
-        slice_1_inputs = [f"{shape_4_name}/output_0", "/model/constants/TensorProto.INT64/1D/-1", f"/model/constants/TensorProto.INT64/1D/{np.iinfo(np.int64).max}", "/model/constants/TensorProto.INT64/1D/0"]
+        slice_1_inputs = [f"{shape_4_name}/output_0", "/model/constants/TensorProto.INT64/1D/-1", f"/model/constants/TensorProto.INT64/1D/{torch.iinfo(torch.int64).max}", "/model/constants/TensorProto.INT64/1D/0"]
         self.make_slice(slice_1_name, slice_1_inputs, dtype=TensorProto.INT64, shape=[1])
         squeeze_1_name = f"{basename}/Squeeze_1"
         squeeze_1_inputs = [f"{slice_1_name}/output_0", "/model/constants/TensorProto.INT64/1D/0"]
@@ -2435,7 +2471,7 @@ def make_input_ids_subgraph(self, basename, past_key_gather_name):
         shape_5_name = f"{basename}/Shape_5"
         self.make_shape(shape_5_name, f"{constant_shape_name}/output_0", shape=[2])
         slice_2_name = f"{basename}/Slice_2"
-        slice_2_inputs = [f"{shape_5_name}/output_0", "/model/constants/TensorProto.INT64/1D/-1", f"/model/constants/TensorProto.INT64/1D/{np.iinfo(np.int64).max}", "/model/constants/TensorProto.INT64/1D/0"]
+        slice_2_inputs = [f"{shape_5_name}/output_0", "/model/constants/TensorProto.INT64/1D/-1", f"/model/constants/TensorProto.INT64/1D/{torch.iinfo(torch.int64).max}", "/model/constants/TensorProto.INT64/1D/0"]
         self.make_slice(slice_2_name, slice_2_inputs, dtype=TensorProto.INT64, shape=[1])
         squeeze_2_name = f"{basename}/Squeeze_2"
         squeeze_2_inputs = [f"{slice_2_name}/output_0", "/model/constants/TensorProto.INT64/1D/0"]
@@ -2496,7 +2532,7 @@ def make_attention_mask_subgraph(self, basename, unsqueeze_for_concat):
         cast_2_name = f"{basename}/Cast_2"
         self.make_cast(cast_2_name, f"{sub_name}/output_0", dtype=TensorProto.BOOL, shape=["unk", "unk", "unk", "unk"])
         where_2_name = f"{basename}/Where_2"
-        where_2_inputs = [f"{cast_2_name}/output_0", f"/model/constants/{self.to_str_dtype[self.io_dtype]}/0D/{np.finfo(self.to_numpy_dtype[self.io_dtype]).min}", f"{sub_name}/output_0"]
+        where_2_inputs = [f"{cast_2_name}/output_0", f"/model/constants/{self.to_str_dtype[self.io_dtype]}/0D/{torch.finfo(self.to_torch_dtype[self.io_dtype]).min}", f"{sub_name}/output_0"]
         self.make_where(where_2_name, where_2_inputs, dtype=self.io_dtype, shape=["unk", "unk", "unk", "unk"])
 
         return where_2_name
@@ -2852,10 +2888,10 @@ def make_rotary_embedding_caches(self, **kwargs):
 
             # Slice cos/sin caches from (M, H) to (M, H/2)
             hidden_dim = cos_cache.shape[-1]
-            cos_cache = cos_cache.squeeze()[:, : (hidden_dim // 2)].detach().numpy()
-            cos_cache = cos_cache.astype(self.to_numpy_dtype[self.io_dtype])
-            sin_cache = sin_cache.squeeze()[:, : (hidden_dim // 2)].detach().numpy()
-            sin_cache = sin_cache.astype(self.to_numpy_dtype[self.io_dtype])
+            cos_cache = cos_cache.squeeze()[:, : (hidden_dim // 2)].detach()
+            cos_cache = cos_cache.to(self.to_torch_dtype[self.io_dtype])
+            sin_cache = sin_cache.squeeze()[:, : (hidden_dim // 2)].detach()
+            sin_cache = sin_cache.to(self.to_torch_dtype[self.io_dtype])
 
             # Slice cos/sin caches from (M, H/2) to (M, R/2) if partial rotary embeddings are used
             if self.rotemb_attrs["partial_rotary_factor"] != 1.0:
@@ -2864,8 +2900,8 @@ def make_rotary_embedding_caches(self, **kwargs):
 
             if self.rotemb_attrs["save_caches"]:
                 # Save cos/sin caches to disk
-                self.make_external_tensor(cos_cache, cos_cache_name)
-                self.make_external_tensor(sin_cache, sin_cache_name)
+                self.make_external_tensor(cos_cache.contiguous(), cos_cache_name)
+                self.make_external_tensor(sin_cache.contiguous(), sin_cache_name)
             else:
                 # Return cos/sin caches since they will be custom-saved
                 return cos_cache, sin_cache
@@ -2945,11 +2981,11 @@ def calculate_block_mask(self):
 
         # Create tensors for row indices and col indices
         crows_name = "block_row_indices"
-        self.make_external_tensor(crows.detach().numpy().astype(np.int32), crows_name)
+        self.make_external_tensor(crows.detach().to(torch.int32).contiguous(), crows_name)
         self.mask_attrs["block_row_indices"] = crows_name
 
         cols_name = "block_col_indices"
-        self.make_external_tensor(cols.detach().numpy().astype(np.int32), cols_name)
+        self.make_external_tensor(cols.detach().to(torch.int32).contiguous(), cols_name)
         self.mask_attrs["block_col_indices"] = cols_name
 
     def make_attention(self, layer_id, attention, root_input, **kwargs):
@@ -3022,11 +3058,11 @@ def make_mlp_proj(self, layer_id, mlp, root_input):
         up_matmul_name = f"/model/layers.{layer_id}/mlp/up_proj/MatMul"
         self.make_matmul(mlp.up_proj, up_matmul_name, root_input)
         up_add_name = f"/model/layers.{layer_id}/mlp/up_proj/Add"
-        self.make_add_bias(mlp.up_proj.bias.detach().numpy(), up_add_name, f"{up_matmul_name}/output_0")
+        self.make_add_bias(mlp.up_proj.bias.detach(), up_add_name, f"{up_matmul_name}/output_0")
 
         # Left path
         slice_1_name = f"/model/layers.{layer_id}/mlp/gelu/Slice"
-        slice_1_inputs = [f"{up_add_name}/output_0", "/model/constants/TensorProto.INT64/1D/0", f"/model/constants/TensorProto.INT64/1D/{np.iinfo(np.int64).max}", "/model/constants/TensorProto.INT64/1D/-1", "/model/constants/TensorProto.INT64/1D/2"]
+        slice_1_inputs = [f"{up_add_name}/output_0", "/model/constants/TensorProto.INT64/1D/0", f"/model/constants/TensorProto.INT64/1D/{torch.iinfo(torch.int64).max}", "/model/constants/TensorProto.INT64/1D/-1", "/model/constants/TensorProto.INT64/1D/2"]
         self.make_slice(slice_1_name, slice_1_inputs, dtype=self.io_dtype, shape=["batch_size", "sequence_length", self.intermediate_size])
         cast_1_name = f"/model/layers.{layer_id}/mlp/gelu/Cast"
         self.make_cast(cast_1_name, f"{slice_1_name}/output_0", dtype=TensorProto.FLOAT, shape=["batch_size", "sequence_length", self.intermediate_size])
@@ -3043,7 +3079,7 @@ def make_mlp_proj(self, layer_id, mlp, root_input):
 
         # Right path
         slice_2_name = f"/model/layers.{layer_id}/mlp/linear/Slice"
-        slice_2_inputs = [f"{up_add_name}/output_0", "/model/constants/TensorProto.INT64/1D/1", f"/model/constants/TensorProto.INT64/1D/{np.iinfo(np.int64).max}", "/model/constants/TensorProto.INT64/1D/-1", "/model/constants/TensorProto.INT64/1D/2"]
+        slice_2_inputs = [f"{up_add_name}/output_0", "/model/constants/TensorProto.INT64/1D/1", f"/model/constants/TensorProto.INT64/1D/{torch.iinfo(torch.int64).max}", "/model/constants/TensorProto.INT64/1D/-1", "/model/constants/TensorProto.INT64/1D/2"]
         self.make_slice(slice_2_name, slice_2_inputs, dtype=self.io_dtype, shape=["batch_size", "sequence_length", self.intermediate_size])
         cast_2_name = f"/model/layers.{layer_id}/mlp/linear/Cast"
         self.make_cast(cast_2_name, f"{slice_2_name}/output_0", dtype=TensorProto.FLOAT, shape=["batch_size", "sequence_length", self.intermediate_size])
@@ -3068,7 +3104,7 @@ def make_mlp_proj(self, layer_id, mlp, root_input):
         down_matmul_name = f"/model/layers.{layer_id}/mlp/down_proj/MatMul"
         self.make_matmul(mlp.down_proj, down_matmul_name, f"{mul_name}/output_0")
         down_add_name = f"/model/layers.{layer_id}/mlp/down_proj/Add"
-        self.make_add_bias(mlp.down_proj.bias.detach().numpy(), down_add_name, f"{down_matmul_name}/output_0")
+        self.make_add_bias(mlp.down_proj.bias.detach(), down_add_name, f"{down_matmul_name}/output_0")
 
         # Assign output 0 of previous MatMul as skip input to next SkipLayerNorm
         self.layernorm_attrs["skip_input"] = f"{down_add_name}/output_0"
@@ -3321,6 +3357,20 @@ def parse_hf_token(hf_token):
     return hf_token
 
 
+def set_io_dtype(precision, execution_provider, extra_options):
+    use_webgpu_fp32 = extra_options.get("use_webgpu_fp32", "0") == "1"
+    if precision in {"int8", "fp32"} or (precision == "int4" and execution_provider == "cpu") or use_webgpu_fp32:
+        # FP32 precision
+        return TensorProto.FLOAT
+
+    if precision == "bf16":
+        # BF16 precision
+        return TensorProto.BFLOAT16
+
+    # FP16 precision
+    return TensorProto.FLOAT16
+
+
 def create_model(model_name, input_path, output_dir, precision, execution_provider, cache_dir, **extra_options):
     # Create cache and output directories
     os.makedirs(output_dir, exist_ok=True)
@@ -3338,8 +3388,7 @@ def create_model(model_name, input_path, output_dir, precision, execution_provid
         config.update(peft_config.__dict__)
 
     # Set input/output precision of ONNX model
-    use_webgpu_fp32 = extra_options.get("use_webgpu_fp32", "0") == "1"
-    io_dtype = TensorProto.FLOAT if precision in {"int8", "fp32"} or (precision == "int4" and execution_provider == "cpu") or use_webgpu_fp32 else TensorProto.FLOAT16
+    io_dtype = set_io_dtype(precision, execution_provider, extra_options)
 
     if "config_only" not in extra_options:
         # List architecture options in alphabetical order
@@ -3450,7 +3499,7 @@ def get_args():
         "-p",
         "--precision",
         required=True,
-        choices=["int4", "fp16", "fp32"],
+        choices=["int4", "bf16", "fp16", "fp32"],
         help="Precision of model",
     )
 
@@ -3527,7 +3576,7 @@ def get_args():
     )
 
     args = parser.parse_args()
-    print("Valid precision + execution provider combinations are: FP32 CPU, FP32 CUDA, FP16 CUDA, FP16 DML, INT4 CPU, INT4 CUDA, INT4 DML, INT4 WEBGPU")
+    print("Valid precision + execution provider combinations are: FP32 CPU, FP32 CUDA, FP16 CUDA, FP16 DML, BF16 CPU, BF16 CUDA, INT4 CPU, INT4 CUDA, INT4 DML, INT4 WEBGPU")
     return args
 
 if __name__ == '__main__':
diff --git a/src/python/py/models/quantized_model.py b/src/python/py/models/quantized_model.py
index 3a7b3f46c4..e3af11beda 100644
--- a/src/python/py/models/quantized_model.py
+++ b/src/python/py/models/quantized_model.py
@@ -112,9 +112,6 @@ def __init__(self, quant_type, input_path, quant_attrs, q_size, kv_size, interme
                     local_bits = self.get_layer_bits(name)  # codeql[py/init-calls-subclass]
                     local_group_size = self.get_layer_group_size(name)  # codeql[py/init-calls-subclass]
 
-                    if tensor.dtype == torch.bfloat16:
-                        # Cast bfloat16 to float32 since NumPy does not support bfloat16
-                        tensor = tensor.to(torch.float32)
                     if name == "model.embed_tokens.weight" or name == "transformer.embedding.word_embeddings.weight":
                         self.embedding.weight = tensor
                     elif name == "model.norm.weight" or name == "transformer.encoder.final_layernorm.weight":

From 5fa15f2d1dd2b62226e973a22fab65ca4a8eada7 Mon Sep 17 00:00:00 2001
From: Kunal Vaishnavi <kvaishnavi@microsoft.com>
Date: Mon, 21 Apr 2025 18:43:30 +0000
Subject: [PATCH 02/26] Refactor multi-cache calculation

---
 src/python/py/models/builder.py | 142 ++++++++++++--------------------
 1 file changed, 51 insertions(+), 91 deletions(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index 8e4cbcc6bd..c8bc3bb40c 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -523,7 +523,7 @@ def order_repeated_field(self, repeated_proto, key_name, order):
         order = list(order)
         repeated_proto.sort(key=lambda x: order.index(getattr(x, key_name)))
 
-    def make_external_tensor(self, tensor, name, **kwargs):
+    def make_tensor_proto_from_tensor(self, tensor, name):
         raw_data = bytes(
             (ctypes.c_ubyte * tensor.element_size() * tensor.numel()).from_address(
                 tensor.data_ptr()
@@ -536,6 +536,10 @@ def make_external_tensor(self, tensor, name, **kwargs):
             vals=raw_data,
             raw=True,
         )
+        return tensor_proto
+
+    def make_external_tensor(self, tensor, name, **kwargs):
+        tensor_proto = self.make_tensor_proto_from_tensor(tensor, name)
 
         filename = f"{name}.bin"
         external_data_helper.set_external_data(tensor_proto, location=filename)
@@ -544,7 +548,6 @@ def make_external_tensor(self, tensor, name, **kwargs):
         tensor_proto.ClearField("raw_data")
         tensor_proto.data_location = TensorProto.EXTERNAL
 
-        # TODO: is this still needed for DML?
         if kwargs.get("unpack_int4", False) and self.onnx_dtype == 'int4':
             tensor_proto.data_type = TensorProto.UINT4
             tensor_proto.dims[-1] *= 2
@@ -626,18 +629,7 @@ def make_constant(self, name):
 
         torch_dtype = self.to_torch_dtype[onnx_dtype]
         tensor = torch.tensor(num if dims == "0D" else list(num) if type(num) == tuple else [num], dtype=torch_dtype).contiguous()
-        raw_data = bytes(
-            (ctypes.c_ubyte * tensor.element_size() * tensor.numel()).from_address(
-                tensor.data_ptr()
-            )
-        )
-        value = helper.make_tensor(
-            name=name.replace("constants", "torch_helper"),
-            data_type=torch.onnx._type_utils.JitScalarType.from_dtype(tensor.dtype).onnx_type(),
-            dims=tensor.shape,
-            vals=raw_data,
-            raw=True,
-        )
+        value = self.make_tensor_proto_from_tensor(tensor, name.replace("constants", "torch_helper"))
 
         node_name = name.replace("constants", "constant_nodes")
         self.make_node("Constant", inputs=[], outputs=[name], name=node_name, value=value)
@@ -846,7 +838,6 @@ def make_matmul_int4(self, matmul, basename, root_input, **kwargs):
 
         return name
 
-    # TODO: are the booleans in make_external_tensor still needed?
     def make_dequantize_linear(self, dequantize_name, quantized_op):
         # Input weights are quantized, save quantized MatMul weights for onnx model
         qweight = dequantize_name[1:].replace("/", ".") + ".qweight"
@@ -1195,34 +1186,50 @@ def make_rotary_embedding(self, name, root_input, **kwargs):
         )
         self.make_value_info(output, self.io_dtype, shape=['batch_size', 'sequence_length', self.head_size * num_heads])
 
-    def make_rotary_embedding_multi_cache(self):
-        if_cos_cache_output, if_sin_cache_output = "cos_cache", "sin_cache"
+    def make_rotary_embedding_multi_cache(self, **kwargs):
+        cos_cache_name = kwargs.get("cos_cache_name", "cos_cache")
+        sin_cache_name = kwargs.get("sin_cache_name", "sin_cache")
 
-        # Create caches for when sequence_length > self.original_context_length
+        # Set cache attributes for when sequence_length > self.original_context_length
         self.rotemb_attrs["rescale_factors"] = self.rotemb_attrs["multi_cache"]["long_factor"]
         self.rotemb_attrs["cache_length"] = self.context_length
         self.rotemb_attrs["mscale"] = self.rotemb_attrs["multi_cache"]["long_mscale"]
 
-        # DML doesn't support dynamic selection of the cos/sin cache, so we always use the biggest one
-        if self.ep == "dml":
-            self.make_rotary_embedding_caches()
-            self.make_value_info(if_cos_cache_output, self.io_dtype, shape=["max_sequence_length", "head_dim / 2"])
-            self.make_value_info(if_sin_cache_output, self.io_dtype, shape=["max_sequence_length", "head_dim / 2"])
-            return
-
-        self.rotemb_attrs["save_caches"] = False
-
+        # Create caches for when sequence_length > self.original_context_length
         cos_cache_large_name, sin_cache_large_name = "cos_cache_large", "sin_cache_large"
-        cos_cache_large, sin_cache_large = self.make_rotary_embedding_caches(cos_cache_name=cos_cache_large_name, sin_cache_name=sin_cache_large_name)
+        if self.ep == "dml":
+            self.rotemb_attrs["save_caches"] = False
+            cos_cache_large, sin_cache_large = self.make_rotary_embedding_caches(cos_cache_name=cos_cache_large_name, sin_cache_name=sin_cache_large_name)
+        else:
+            self.make_rotary_embedding_caches(cos_cache_name=cos_cache_large_name, sin_cache_name=sin_cache_large_name)
+            cos_cache_large = list(filter(lambda i: i.name == cos_cache_large_name, self.initializers))[0]
+            sin_cache_large = list(filter(lambda i: i.name == sin_cache_large_name, self.initializers))[0]
 
-        # Create caches for when sequence_length <= self.original_context_length
+        # Set cache attributes for when sequence_length <= self.original_context_length
         self.rotemb_attrs["rescale_factors"] = self.rotemb_attrs["multi_cache"]["short_factor"]
         self.rotemb_attrs["cache_length"] = self.original_context_length
         self.rotemb_attrs["mscale"] = self.rotemb_attrs["multi_cache"]["short_mscale"]
+        self.rotemb_attrs["create_caches"] = True
+
+        # Create caches for when sequence_length <= self.original_context_length
         cos_cache_small_name, sin_cache_small_name = "cos_cache_small", "sin_cache_small"
-        cos_cache_small, sin_cache_small = self.make_rotary_embedding_caches(cos_cache_name=cos_cache_small_name, sin_cache_name=sin_cache_small_name)
+        if self.ep == "dml":
+            self.rotemb_attrs["save_caches"] = False
+            cos_cache_small, sin_cache_small = self.make_rotary_embedding_caches(cos_cache_name=cos_cache_small_name, sin_cache_name=sin_cache_small_name)
+        else:
+            self.make_rotary_embedding_caches(cos_cache_name=cos_cache_small_name, sin_cache_name=sin_cache_small_name)
+            cos_cache_small = list(filter(lambda i: i.name == cos_cache_small_name, self.initializers))[0]
+            sin_cache_small = list(filter(lambda i: i.name == sin_cache_small_name, self.initializers))[0]
 
-        self.rotemb_attrs["create_caches"] = False
+        if self.ep == "dml":
+            # Concat small and large cos/sin caches for DML EP only
+            cos_cache = torch.cat((cos_cache_small, cos_cache_large), dim=0)
+            sin_cache = torch.cat((sin_cache_small, sin_cache_large), dim=0)
+            # Save cos/sin caches to disk
+            self.make_external_tensor(cos_cache.contiguous(), cos_cache_name)
+            self.make_external_tensor(sin_cache.contiguous(), sin_cache_name)
+            # Do NOT make the subgraph with the If node for DML EP.
+            return
 
         # Make the following subgraph to decide which cos/sin caches to use in the rotary embeddings
         #
@@ -1242,40 +1249,38 @@ def make_rotary_embedding_multi_cache(self):
         self.make_greater(greater_name, greater_inputs, shape=[])
         if_name = f"{basename}/If"
         self.make_node(
-            "If", inputs=[f"{greater_name}/output_0"], outputs=[if_cos_cache_output, if_sin_cache_output], name=if_name,
+            "If", inputs=[f"{greater_name}/output_0"], outputs=[cos_cache_name, sin_cache_name], name=if_name,
             then_branch=self.make_graph(
                 name="large_rotemb_caches_graph",
                 inputs=[],
                 outputs=[
-                    helper.make_tensor_value_info(cos_cache_large_name, self.io_dtype, shape=cos_cache_large.shape),
-                    helper.make_tensor_value_info(sin_cache_large_name, self.io_dtype, shape=sin_cache_large.shape),
+                    helper.make_tensor_value_info(cos_cache_large_name, self.io_dtype, shape=cos_cache_large.dims),
+                    helper.make_tensor_value_info(sin_cache_large_name, self.io_dtype, shape=sin_cache_large.dims),
                 ],
                 initializer=[],
                 value_info=[],
                 nodes=[
-                    # TODO: fix constant node values since caches can be bf16
-                    helper.make_node("Constant", inputs=[], outputs=[cos_cache_large_name], name="/large/cos_cache/Constant", value=numpy_helper.from_array(cos_cache_large)),
-                    helper.make_node("Constant", inputs=[], outputs=[sin_cache_large_name], name="/large/sin_cache/Constant", value=numpy_helper.from_array(sin_cache_large)),
+                    helper.make_node("Constant", inputs=[], outputs=[cos_cache_large_name], name="/large/cos_cache/Constant", value=cos_cache_large),
+                    helper.make_node("Constant", inputs=[], outputs=[sin_cache_large_name], name="/large/sin_cache/Constant", value=sin_cache_large),
                 ],
             ),
             else_branch=self.make_graph(
                 name="small_rotemb_caches_graph",
                 inputs=[],
                 outputs=[
-                    helper.make_tensor_value_info(cos_cache_small_name, self.io_dtype, shape=cos_cache_small.shape),
-                    helper.make_tensor_value_info(sin_cache_small_name, self.io_dtype, shape=sin_cache_small.shape),
+                    helper.make_tensor_value_info(cos_cache_small_name, self.io_dtype, shape=cos_cache_small.dims),
+                    helper.make_tensor_value_info(sin_cache_small_name, self.io_dtype, shape=sin_cache_small.dims),
                 ],
                 initializer=[],
                 value_info=[],
                 nodes=[
-                    # TODO: fix constant node values since caches can be bf16
-                    helper.make_node("Constant", inputs=[], outputs=[cos_cache_small_name], name="/small/cos_cache/Constant", value=numpy_helper.from_array(cos_cache_small)),
-                    helper.make_node("Constant", inputs=[], outputs=[sin_cache_small_name], name="/small/sin_cache/Constant", value=numpy_helper.from_array(sin_cache_small)),
+                    helper.make_node("Constant", inputs=[], outputs=[cos_cache_small_name], name="/small/cos_cache/Constant", value=cos_cache_small),
+                    helper.make_node("Constant", inputs=[], outputs=[sin_cache_small_name], name="/small/sin_cache/Constant", value=sin_cache_small),
                 ],
             ),
         )
-        self.make_value_info(if_cos_cache_output, self.io_dtype, shape=["max_sequence_length", "head_dim / 2"])
-        self.make_value_info(if_sin_cache_output, self.io_dtype, shape=["max_sequence_length", "head_dim / 2"])
+        self.make_value_info(cos_cache_name, self.io_dtype, shape=["max_sequence_length", "head_dim / 2"])
+        self.make_value_info(sin_cache_name, self.io_dtype, shape=["max_sequence_length", "head_dim / 2"])
 
     def make_qk_norm(self, layer_id, attention):
         # Make subgraph to compute SimplifiedLayerNorm after Q and K MatMuls in attention:
@@ -2445,10 +2450,9 @@ def make_input_ids_subgraph(self, basename, past_key_gather_name):
         concat_2_name = f"{basename}/Concat_2"
         concat_inputs = [f"{unsqueeze_4_name}/output_0", f"{unsqueeze_5_name}/output_0"]
         self.make_concat(concat_2_name, concat_inputs, dtype=TensorProto.INT64, shape=[2], axis=0)
-        # TODO: use to_torch_dtype for constant node
         constant_shape_name = f"{basename}/ConstantOfShape_2"
-        constant_shape_numpy_dtype = self.to_numpy_dtype[self.io_dtype]
-        constant_shape_value = numpy_helper.from_array(np.array([np.finfo(constant_shape_numpy_dtype).min], dtype=constant_shape_numpy_dtype))
+        constant_shape_torch_dtype = self.to_torch_dtype[self.io_dtype]
+        constant_shape_value = self.make_tensor_proto_from_tensor(torch.tensor([torch.finfo(constant_shape_torch_dtype).min], dtype=constant_shape_torch_dtype))
         self.make_constant_of_shape(constant_shape_name, f"{concat_2_name}/output_0", value=constant_shape_value, dtype=self.io_dtype, shape=['unk', 'unk'])
 
         # Top path
@@ -2865,50 +2869,6 @@ def make_position_ids_reformatting(self):
         self.make_add(add_1_name, add_1_inputs, dtype=proto_dtype, shape=["batch_size", "sequence_length"])
 
         return add_1_name
-        
-    def make_rotary_embedding_caches(self, **kwargs):
-        if self.ep != "dml":
-            cos_cache_name, sin_cache_name = super().make_rotary_embedding_caches(**kwargs)
-            return cos_cache_name, sin_cache_name
-
-        cos_cache_name = kwargs.get("cos_cache_name", "cos_cache")
-        sin_cache_name = kwargs.get("sin_cache_name", "sin_cache")
-
-        if self.rotemb_attrs["create_caches"]:
-            # Create 4K and 128K cos/sin caches
-            cos_cache_large, sin_cache_large = self.make_rotary_embedding_caches_from_scratch()
-            self.rotemb_attrs["rescale_factors"] = self.rotemb_attrs["multi_cache"]["short_factor"]
-            self.rotemb_attrs["cache_length"] = self.original_context_length
-            self.rotemb_attrs["mscale"] = self.rotemb_attrs["multi_cache"]["short_mscale"]
-            cos_cache_small, sin_cache_small = self.make_rotary_embedding_caches_from_scratch()
-
-            # Concat 4K and 128K cos/sin caches for DML EP only
-            cos_cache = torch.cat((cos_cache_small, cos_cache_large), dim=0)
-            sin_cache = torch.cat((sin_cache_small, sin_cache_large), dim=0)
-
-            # Slice cos/sin caches from (M, H) to (M, H/2)
-            hidden_dim = cos_cache.shape[-1]
-            cos_cache = cos_cache.squeeze()[:, : (hidden_dim // 2)].detach()
-            cos_cache = cos_cache.to(self.to_torch_dtype[self.io_dtype])
-            sin_cache = sin_cache.squeeze()[:, : (hidden_dim // 2)].detach()
-            sin_cache = sin_cache.to(self.to_torch_dtype[self.io_dtype])
-
-            # Slice cos/sin caches from (M, H/2) to (M, R/2) if partial rotary embeddings are used
-            if self.rotemb_attrs["partial_rotary_factor"] != 1.0:
-                cos_cache = cos_cache[:, : (self.rotemb_attrs["rotary_embedding_dim"] // 2)]
-                sin_cache = sin_cache[:, : (self.rotemb_attrs["rotary_embedding_dim"] // 2)]
-
-            if self.rotemb_attrs["save_caches"]:
-                # Save cos/sin caches to disk
-                self.make_external_tensor(cos_cache.contiguous(), cos_cache_name)
-                self.make_external_tensor(sin_cache.contiguous(), sin_cache_name)
-            else:
-                # Return cos/sin caches since they will be custom-saved
-                return cos_cache, sin_cache
-
-            self.rotemb_attrs["create_caches"] = False
-
-        return cos_cache_name, sin_cache_name
 
 
 class Phi3SmallModel(Model):

From 084936a3606add4b8f8b9aeab782191cc05c675d Mon Sep 17 00:00:00 2001
From: Kunal Vaishnavi <kvaishnavi@microsoft.com>
Date: Mon, 21 Apr 2025 20:30:05 +0000
Subject: [PATCH 03/26] Fix how caches are saved in if node

---
 src/python/py/models/builder.py | 38 ++++++++++++---------------------
 1 file changed, 14 insertions(+), 24 deletions(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index c8bc3bb40c..192e816c97 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -1161,6 +1161,8 @@ def make_rotary_embedding_caches(self, **kwargs):
                 cos_cache = cos_cache[:, : (self.rotemb_attrs["rotary_embedding_dim"] // 2)]
                 sin_cache = sin_cache[:, : (self.rotemb_attrs["rotary_embedding_dim"] // 2)]
 
+            self.rotemb_attrs["create_caches"] = False
+
             if self.rotemb_attrs["save_caches"]:
                 # Save cos/sin caches to disk
                 self.make_external_tensor(cos_cache.contiguous(), cos_cache_name)
@@ -1169,8 +1171,6 @@ def make_rotary_embedding_caches(self, **kwargs):
                 # Return cos/sin caches since they will be custom-saved
                 return cos_cache, sin_cache
 
-            self.rotemb_attrs["create_caches"] = False
-
         return cos_cache_name, sin_cache_name
 
     def make_rotary_embedding(self, name, root_input, **kwargs):
@@ -1197,13 +1197,8 @@ def make_rotary_embedding_multi_cache(self, **kwargs):
 
         # Create caches for when sequence_length > self.original_context_length
         cos_cache_large_name, sin_cache_large_name = "cos_cache_large", "sin_cache_large"
-        if self.ep == "dml":
-            self.rotemb_attrs["save_caches"] = False
-            cos_cache_large, sin_cache_large = self.make_rotary_embedding_caches(cos_cache_name=cos_cache_large_name, sin_cache_name=sin_cache_large_name)
-        else:
-            self.make_rotary_embedding_caches(cos_cache_name=cos_cache_large_name, sin_cache_name=sin_cache_large_name)
-            cos_cache_large = list(filter(lambda i: i.name == cos_cache_large_name, self.initializers))[0]
-            sin_cache_large = list(filter(lambda i: i.name == sin_cache_large_name, self.initializers))[0]
+        self.rotemb_attrs["save_caches"] = False
+        cos_cache_large, sin_cache_large = self.make_rotary_embedding_caches(cos_cache_name=cos_cache_large_name, sin_cache_name=sin_cache_large_name)
 
         # Set cache attributes for when sequence_length <= self.original_context_length
         self.rotemb_attrs["rescale_factors"] = self.rotemb_attrs["multi_cache"]["short_factor"]
@@ -1213,13 +1208,8 @@ def make_rotary_embedding_multi_cache(self, **kwargs):
 
         # Create caches for when sequence_length <= self.original_context_length
         cos_cache_small_name, sin_cache_small_name = "cos_cache_small", "sin_cache_small"
-        if self.ep == "dml":
-            self.rotemb_attrs["save_caches"] = False
-            cos_cache_small, sin_cache_small = self.make_rotary_embedding_caches(cos_cache_name=cos_cache_small_name, sin_cache_name=sin_cache_small_name)
-        else:
-            self.make_rotary_embedding_caches(cos_cache_name=cos_cache_small_name, sin_cache_name=sin_cache_small_name)
-            cos_cache_small = list(filter(lambda i: i.name == cos_cache_small_name, self.initializers))[0]
-            sin_cache_small = list(filter(lambda i: i.name == sin_cache_small_name, self.initializers))[0]
+        self.rotemb_attrs["save_caches"] = False
+        cos_cache_small, sin_cache_small = self.make_rotary_embedding_caches(cos_cache_name=cos_cache_small_name, sin_cache_name=sin_cache_small_name)
 
         if self.ep == "dml":
             # Concat small and large cos/sin caches for DML EP only
@@ -1254,28 +1244,28 @@ def make_rotary_embedding_multi_cache(self, **kwargs):
                 name="large_rotemb_caches_graph",
                 inputs=[],
                 outputs=[
-                    helper.make_tensor_value_info(cos_cache_large_name, self.io_dtype, shape=cos_cache_large.dims),
-                    helper.make_tensor_value_info(sin_cache_large_name, self.io_dtype, shape=sin_cache_large.dims),
+                    helper.make_tensor_value_info(cos_cache_large_name, self.io_dtype, shape=cos_cache_large.shape),
+                    helper.make_tensor_value_info(sin_cache_large_name, self.io_dtype, shape=sin_cache_large.shape),
                 ],
                 initializer=[],
                 value_info=[],
                 nodes=[
-                    helper.make_node("Constant", inputs=[], outputs=[cos_cache_large_name], name="/large/cos_cache/Constant", value=cos_cache_large),
-                    helper.make_node("Constant", inputs=[], outputs=[sin_cache_large_name], name="/large/sin_cache/Constant", value=sin_cache_large),
+                    helper.make_node("Constant", inputs=[], outputs=[cos_cache_large_name], name="/large/cos_cache/Constant", value=self.make_tensor_proto_from_tensor(cos_cache_large, "cos_cache_large_torch_helper")),
+                    helper.make_node("Constant", inputs=[], outputs=[sin_cache_large_name], name="/large/sin_cache/Constant", value=self.make_tensor_proto_from_tensor(sin_cache_large, "sin_cache_large_torch_helper")),
                 ],
             ),
             else_branch=self.make_graph(
                 name="small_rotemb_caches_graph",
                 inputs=[],
                 outputs=[
-                    helper.make_tensor_value_info(cos_cache_small_name, self.io_dtype, shape=cos_cache_small.dims),
-                    helper.make_tensor_value_info(sin_cache_small_name, self.io_dtype, shape=sin_cache_small.dims),
+                    helper.make_tensor_value_info(cos_cache_small_name, self.io_dtype, shape=cos_cache_small.shape),
+                    helper.make_tensor_value_info(sin_cache_small_name, self.io_dtype, shape=sin_cache_small.shape),
                 ],
                 initializer=[],
                 value_info=[],
                 nodes=[
-                    helper.make_node("Constant", inputs=[], outputs=[cos_cache_small_name], name="/small/cos_cache/Constant", value=cos_cache_small),
-                    helper.make_node("Constant", inputs=[], outputs=[sin_cache_small_name], name="/small/sin_cache/Constant", value=sin_cache_small),
+                    helper.make_node("Constant", inputs=[], outputs=[cos_cache_small_name], name="/small/cos_cache/Constant", value=self.make_tensor_proto_from_tensor(cos_cache_small, "cos_cache_small_torch_helper")),
+                    helper.make_node("Constant", inputs=[], outputs=[sin_cache_small_name], name="/small/sin_cache/Constant", value=self.make_tensor_proto_from_tensor(sin_cache_small, "sin_cache_small_torch_helper")),
                 ],
             ),
         )

From eb6f0fbdd72bc8c7bfb659cb5bf479522befb36a Mon Sep 17 00:00:00 2001
From: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Date: Wed, 23 Apr 2025 17:10:16 -0700
Subject: [PATCH 04/26] Fix node names for logit softcapping

---
 src/python/py/models/builder.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index 14e7e445dc..eb51050b8d 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -2143,14 +2143,14 @@ def make_lm_head(self, lm_head):
 
         if softcap_exists:
             # Add final logit softcapping (Div --> Tanh --> Mul)
-            div_name = "/lm_head/Div"
+            div_name = "/lm_head/softcap/Div"
             div_inputs = [f"{lm_name}/output_0", f"/model/constants/{self.to_str_dtype[self.io_dtype]}/0D/{self.lm_head_attrs['softcap']}"]
             self.make_div(div_name, div_inputs, dtype=self.io_dtype, shape=['batch_size', 'sequence_length', self.vocab_size])
 
-            tanh_name = "/lm_head/Tanh"
+            tanh_name = "/lm_head/softcap/Tanh"
             self.make_tanh(tanh_name, f"{div_name}/output_0", dtype=self.io_dtype, shape=['batch_size', 'sequence_length', self.vocab_size])
 
-            mul_name = "/lm_head/Mul"
+            mul_name = "/lm_head/softcap/Mul"
             mul_inputs = [f"{tanh_name}/output_0", f"/model/constants/{self.to_str_dtype[self.io_dtype]}/0D/{self.lm_head_attrs['softcap']}"]
             mul_output = "logits"
             self.make_node('Mul', inputs=mul_inputs, outputs=[mul_output], name=mul_name)

From d3b5ee089f558281f63fd374104dfa9184b6eefc Mon Sep 17 00:00:00 2001
From: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Date: Wed, 23 Apr 2025 17:37:15 -0700
Subject: [PATCH 05/26] Remove extra multiply

---
 src/python/py/models/builder.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index eb51050b8d..a30f5f2aa5 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -2754,7 +2754,6 @@ class Gemma2Model(GemmaModel):
     def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
         super().__init__(config, io_dtype, onnx_dtype, ep, cache_dir, extra_options)
         self.attention_attrs["scale"] = config.query_pre_attn_scalar ** -0.5
-        self.lm_head_attrs["scale"] = config.final_logit_softcapping if config.final_logit_softcapping is not None else 1.0
         self.is_local = lambda layer_id: layer_id % 2 == 1
 
     def make_layer(self, layer_id, layer):

From 4880922d57f1ac415a5ff900be1786cc536a7436 Mon Sep 17 00:00:00 2001
From: Kunal Vaishnavi <kvaishnavi@microsoft.com>
Date: Thu, 24 Apr 2025 17:55:11 +0000
Subject: [PATCH 06/26] Add missing final norm for Gemma-3 multimodal

---
 src/python/py/models/builder.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index a30f5f2aa5..77f480426a 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -2251,9 +2251,10 @@ def has_final_norm(self, module, model):
        hf_norm = hasattr(model, "model") and hasattr(model.model, "norm") and module == model.model.norm
        hf_final_layernorm = hasattr(model, "model") and hasattr(model.model, "final_layernorm") and module == model.model.final_layernorm
        hf_transformer_final_layernorm = hasattr(model, "transformer") and hasattr(model.transformer, "encoder") and hasattr(model.transformer.encoder, "final_layernorm") and module == model.transformer.encoder.final_layernorm
+       hf_multimodal_final_layernorm = hasattr(model, "language_model") and hasattr(model.language_model, "model") and hasattr(model.language_model.model, "norm") and module == model.language_model.model.norm
        # GGUF names
        gguf_final_norm = hasattr(model, "final_norm") and module == model.final_norm
-       return hf_norm or hf_final_layernorm or hf_transformer_final_layernorm or gguf_final_norm
+       return hf_norm or hf_final_layernorm or hf_transformer_final_layernorm or hf_multimodal_final_layernorm or gguf_final_norm
 
     def make_preprocessing_nodes(self):
         self.make_attention_mask_reformatting()

From 2e00e80cc3499f530379d238dbe99d8300e2b453 Mon Sep 17 00:00:00 2001
From: Kunal Vaishnavi <kvaishnavi@microsoft.com>
Date: Thu, 24 Apr 2025 21:42:31 +0000
Subject: [PATCH 07/26] Fix naming of model type

---
 src/python/py/models/builder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index 77f480426a..92053db93e 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -367,7 +367,7 @@ def make_genai_config(self, model_name_or_path, extra_kwargs, out_dir):
                 },
                 "eos_token_id": config.eos_token_id,
                 "pad_token_id": config.pad_token_id if hasattr(config, "pad_token_id") and config.pad_token_id is not None else config.eos_token_id[0] if isinstance(config.eos_token_id, list) else config.eos_token_id,
-                "type": self.model_type[ : self.model_type.find("For")].lower(),
+                "type": self.model_type[ : self.model_type.find("For") if "For" in self.model_type else len(self.model_type)].lower(),
                 "vocab_size": self.vocab_size,
             },
             "search": {

From 196c1c2ace6b7b475be1852a53199dcbed159d64 Mon Sep 17 00:00:00 2001
From: Kunal Vaishnavi <kvaishnavi@microsoft.com>
Date: Thu, 24 Apr 2025 22:03:41 +0000
Subject: [PATCH 08/26] Add Gemma-3 system prompt to chat example

---
 examples/python/model-chat.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/examples/python/model-chat.py b/examples/python/model-chat.py
index 1f3bc005a6..35f40080fa 100644
--- a/examples/python/model-chat.py
+++ b/examples/python/model-chat.py
@@ -55,6 +55,8 @@ def main(args):
             print("Using Chat Template for LLAMA 3, if you are using LLAMA  2 please pass the argument --chat_template '{input} [/INST]')")
         elif model_type.startswith("qwen2"):
             args.chat_template = '<|im_start|>user\n{input}<|im_end|>\n<|im_start|>assistant\n'
+        elif model_type == "gemma3_text":
+            args.chat_template = '<start_of_turn>user\n{system_prompt}{input}<end_of_turn>\n<start_of_turn>model\n'
         else:
             raise ValueError(f"Chat Template for model type {model_type} is not known. Please provide chat template using --chat_template")
 
@@ -81,6 +83,8 @@ def main(args):
             print("Using System Prompt for LLAMA 3, if you are using LLAMA  2 please pass the argument --system_prompt '<s>[INST] <<SYS>>\\n{args.system_prompt}\\n<</SYS>>')")
         elif model_type.startswith("qwen2"):
             system_prompt = f"<|im_start|>system\n{args.system_prompt}<|im_end|>\n"
+        elif model_type == "gemma3_text":
+            system_prompt = f"{args.system_prompt}"
         else:
             system_prompt = args.system_prompt
 
@@ -100,7 +104,7 @@ def main(args):
 
         if args.timings: started_timestamp = time.time()
 
-        prompt = f'{args.chat_template.format(input=text)}'
+        prompt = f'{args.chat_template.format(system_prompt=system_prompt, input=text)}'
         input_tokens = tokenizer.encode(prompt)
         
         generator.append_tokens(input_tokens)

From 7d2fc556a6acc18210f4ad9ef7a6d1a1f47498d7 Mon Sep 17 00:00:00 2001
From: Kunal Vaishnavi <kvaishnavi@microsoft.com>
Date: Fri, 25 Apr 2025 03:51:13 +0000
Subject: [PATCH 09/26] Generate Gemma-3 4B text-only model

---
 src/python/py/models/builder.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index 339c9b1056..7bd3f6cf0e 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -2247,7 +2247,7 @@ def make_model(self, input_path):
         for module in model.modules():
             if isinstance(module, torch.nn.Embedding) or (hasattr(model, "embedding") and module == model.embedding):
                 # Checks (Hugging Face logic) or (GGUF logic)
-                if not self.exclude_embeds:
+                if not self.exclude_embeds and module == model.language_model.model.embed_tokens:
                     # Embedding layer
                     print("Reading embedding layer")
                     self.make_embedding(module.weight.detach())
@@ -3359,8 +3359,9 @@ def create_model(model_name, input_path, output_dir, precision, execution_provid
             for key in text_config:
                 if not hasattr(config, key):
                     setattr(config, key, getattr(text_config, key))
-            extra_options["exclude_embeds"] = True
+            # extra_options["exclude_embeds"] = True
             onnx_model = Gemma3Model(config, io_dtype, precision, execution_provider, cache_dir, extra_options)
+            onnx_model.model_type = "gemma3_text"
         elif config.architectures[0] == "GraniteForCausalLM":
             onnx_model = GraniteModel(config, io_dtype, precision, execution_provider, cache_dir, extra_options)
         elif config.architectures[0] == "LlamaForCausalLM":

From 89040f6d64385399ef0dfe90d1285481b13c2129 Mon Sep 17 00:00:00 2001
From: Kunal Vaishnavi <kvaishnavi@microsoft.com>
Date: Fri, 25 Apr 2025 21:03:46 +0000
Subject: [PATCH 10/26] Cast MatMul from FP32 to BF16

---
 src/python/py/models/builder.py | 22 +++++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index 7bd3f6cf0e..6e10979602 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -792,13 +792,25 @@ def make_matmul_op(self, matmul, basename, root_input, **kwargs):
 
     def make_matmul_float(self, matmul, name, root_input, **kwargs):
         weight = name[1:].replace("/", ".") + ".weight"
-        self.make_external_tensor(matmul.weight.detach().T.to(self.to_torch_dtype[self.io_dtype]).contiguous(), weight)
+        self.make_external_tensor(matmul.weight.detach().T.to(torch.bfloat16).contiguous(), weight)
 
         last_dim = matmul.weight.shape[0]
+        first_dim = matmul.weight.shape[1]
         output = "logits" if kwargs.get("logits", False) else f"{name}/output_0"
-        self.make_node("MatMul", inputs=[root_input, weight], outputs=[output], name=name)
+
+        cast_input = f"{root_input}/Cast"
+        cast_output = f"{output}/Cast"
+
+        if kwargs.get("first_cast", True):
+            self.make_node("Cast", inputs=[root_input], outputs=[cast_input], name=f"{name}/Cast_Input", to=TensorProto.BFLOAT16)
+            self.make_value_info(cast_input, TensorProto.BFLOAT16, shape=["batch_size", "sequence_length", first_dim])
+
+        self.make_node("MatMul", inputs=[cast_input, weight], outputs=[cast_output], name=name)
         self.make_value_info(output, self.io_dtype, shape=['batch_size', 'sequence_length', last_dim])
 
+        self.make_node("Cast", inputs=[cast_output], outputs=[output], name=f"{name}/Cast_Output", to=self.io_dtype)
+        self.make_value_info(cast_output, TensorProto.BFLOAT16, shape=["batch_size", "sequence_length", last_dim])
+
         return name
 
     def make_matmul_int4(self, matmul, basename, root_input, **kwargs):
@@ -1613,10 +1625,10 @@ def make_attention(self, layer_id, attention, root_input, **kwargs):
             q_matmul_name = self.make_matmul(attention.q_proj, q_matmul_basename, root_input)
             self.attention_attrs["q_path"] = f"{q_matmul_name}/output_0"
             k_matmul_basename = f"/model/layers.{layer_id}/attn/k_proj/MatMul"
-            k_matmul_name = self.make_matmul(attention.k_proj, k_matmul_basename, root_input)
+            k_matmul_name = self.make_matmul(attention.k_proj, k_matmul_basename, root_input, first_cast=False)
             self.attention_attrs["k_path"] = f"{k_matmul_name}/output_0"
             v_matmul_basename = f"/model/layers.{layer_id}/attn/v_proj/MatMul"
-            v_matmul_name = self.make_matmul(attention.v_proj, v_matmul_basename, root_input)
+            v_matmul_name = self.make_matmul(attention.v_proj, v_matmul_basename, root_input, first_cast=False)
             self.attention_attrs["v_path"] = f"{v_matmul_name}/output_0"
 
         # Make Add nodes (if bias exists)
@@ -1882,7 +1894,7 @@ def make_mlp_proj(self, layer_id, mlp, root_input):
 
         # Make Up proj nodes
         up_matmul_basename = f"/model/layers.{layer_id}/mlp/up_proj/MatMul"
-        up_matmul_name = self.make_matmul(mlp.up_proj, up_matmul_basename, root_input)
+        up_matmul_name = self.make_matmul(mlp.up_proj, up_matmul_basename, root_input, first_cast=False)
         up_name = up_matmul_name
         if up_bias_exists:
             up_add_name = f"/model/layers.{layer_id}/mlp/up_proj/Add"

From c24d40daf21924015e4045134858f97862b5e741 Mon Sep 17 00:00:00 2001
From: Kunal Vaishnavi <kvaishnavi@microsoft.com>
Date: Mon, 28 Apr 2025 19:35:34 +0000
Subject: [PATCH 11/26] Add missing name for constant node

---
 src/python/py/models/builder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index ad68743eea..6d2f91fdab 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -2462,7 +2462,7 @@ def make_input_ids_subgraph(self, basename, past_key_gather_name):
         self.make_concat(concat_2_name, concat_inputs, dtype=TensorProto.INT64, shape=[2], axis=0)
         constant_shape_name = f"{basename}/ConstantOfShape_2"
         constant_shape_torch_dtype = self.to_torch_dtype[self.io_dtype]
-        constant_shape_value = self.make_tensor_proto_from_tensor(torch.tensor([torch.finfo(constant_shape_torch_dtype).min], dtype=constant_shape_torch_dtype))
+        constant_shape_value = self.make_tensor_proto_from_tensor(torch.tensor([torch.finfo(constant_shape_torch_dtype).min], dtype=constant_shape_torch_dtype), "constant_shape_torch_helper")
         self.make_constant_of_shape(constant_shape_name, f"{concat_2_name}/output_0", value=constant_shape_value, dtype=self.io_dtype, shape=['unk', 'unk'])
 
         # Top path

From b1ec3bd9d21b5a25c130db0625857853bf3e1bc2 Mon Sep 17 00:00:00 2001
From: Nenad Banfic <46795300+nenad1002@users.noreply.github.com>
Date: Tue, 29 Apr 2025 15:10:17 -0400
Subject: [PATCH 12/26] GQA bf16 support (#1429)

---
 src/python/py/models/builder.py | 99 +++++++++++++++++++++++++++++++--
 1 file changed, 93 insertions(+), 6 deletions(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index 6d2f91fdab..a93cc2bcde 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -1171,9 +1171,9 @@ def make_rotary_embedding_caches(self, **kwargs):
             # Slice cos/sin caches from (M, H) to (M, H/2)
             hidden_dim = cos_cache.shape[-1]
             cos_cache = cos_cache.squeeze()[:, : (hidden_dim // 2)].detach()
-            cos_cache = cos_cache.to(self.to_torch_dtype[self.io_dtype])
+            cos_cache = cos_cache.to(self.to_torch_dtype[TensorProto.BFLOAT16])
             sin_cache = sin_cache.squeeze()[:, : (hidden_dim // 2)].detach()
-            sin_cache = sin_cache.to(self.to_torch_dtype[self.io_dtype])
+            sin_cache = sin_cache.to(self.to_torch_dtype[TensorProto.BFLOAT16])
 
             # Slice cos/sin caches from (M, H/2) to (M, R/2) if partial rotary embeddings are used
             if self.rotemb_attrs["partial_rotary_factor"] != 1.0:
@@ -1196,14 +1196,23 @@ def make_rotary_embedding(self, name, root_input, **kwargs):
         cos_cache_name, sin_cache_name = self.make_rotary_embedding_caches()
         num_heads = self.num_kv_heads if "k_rotary" in name else self.num_attn_heads
 
-        inputs = [root_input, kwargs.pop("position_ids"), cos_cache_name, sin_cache_name]
+        casted_root = f"{name}/root_input_bf16"
+        self.make_node(
+            "Cast",
+            inputs=[root_input],
+            outputs=[casted_root],
+            name=f"{name}/Cast_to_bf16",
+            to=TensorProto.BFLOAT16,
+        )
+
+        inputs = [casted_root, kwargs.pop("position_ids"), cos_cache_name, sin_cache_name]
         output = f"{name}/output_0"
         self.make_node(
             "RotaryEmbedding", inputs=inputs, outputs=[output], name=name, domain="com.microsoft",
             interleaved=self.rotemb_attrs["interleaved"], num_heads=(0 if self.rotemb_attrs["partial_rotary_factor"] == 1.0 else num_heads),  # default is 0 in RotaryEmbedding kernel
             rotary_embedding_dim=self.rotemb_attrs["rotary_embedding_dim"],
         )
-        self.make_value_info(output, self.io_dtype, shape=['batch_size', 'sequence_length', self.head_size * num_heads])
+        self.make_value_info(output, TensorProto.BFLOAT16, shape=['batch_size', 'sequence_length', self.head_size * num_heads])
 
     def make_rotary_embedding_multi_cache(self, **kwargs):
         cos_cache_name = kwargs.get("cos_cache_name", "cos_cache")
@@ -1531,7 +1540,7 @@ def make_attention_op(self, name, **kwargs):
         if op_type == "MultiHeadAttention":
             self.make_multi_head_attention(name, add_qk=f"{self.mask_attrs['mask_name']}/output_0", **kwargs)
         elif op_type == "GroupQueryAttention":
-            self.make_group_query_attention(name, seqlens_k=f"{self.mask_attrs['seqlens_k']}/output_0", total_seq_len=f"{self.mask_attrs['total_seq_len']}/output_0", **kwargs)
+            self.make_group_query_attention_with_bf16(name, seqlens_k=f"{self.mask_attrs['seqlens_k']}/output_0", total_seq_len=f"{self.mask_attrs['total_seq_len']}/output_0", **kwargs)
         elif op_type == "SparseAttention":
             self.make_sparse_attention(name, block_row_indices=self.mask_attrs['block_row_indices'], block_col_indices=self.mask_attrs['block_col_indices'], key_total_seq_lens=f"{self.mask_attrs['key_total_seq_lens']}/output_0", total_seq_len=f"{self.mask_attrs['total_seq_len']}/output_0", **kwargs)
         else:
@@ -1567,6 +1576,84 @@ def make_group_query_attention(self, name, **kwargs):
         )
         self.make_value_info(output, self.io_dtype, shape=['batch_size', 'sequence_length', self.head_size * self.num_attn_heads])
 
+    def make_group_query_attention_with_bf16(self, name, **kwargs):
+        raw_inputs = [
+            kwargs["v_path"],
+            kwargs.get("past_k", ""), kwargs.get("past_v", ""),
+            kwargs.get("cos_cache", ""), kwargs.get("sin_cache", ""),
+        ]
+
+        int64_inputs = [kwargs.get("seqlens_k", ""), kwargs.get("total_seq_len", "")]
+
+        bf16_inputs = [kwargs["q_path"], kwargs["k_path"]]
+        for inp in raw_inputs:
+            if inp == "": 
+                bf16_inputs.append("") 
+                continue
+            bf16_inp = f"{inp}_bf16"
+            self.make_node(
+                "Cast",
+                inputs=[inp],
+                outputs=[bf16_inp],
+                name=f"{name}/cast_{inp}_to_bf16",
+                to=TensorProto.BFLOAT16
+            )
+            bf16_inputs.append(bf16_inp)
+
+        bf16_inputs[5:5] = int64_inputs
+
+        out_bf16 = f"{name}/output_0_bf16"
+        pk_bf16  = f"{kwargs['present_k']}_bf16" if kwargs.get("present_k") else ""
+        pv_bf16  = f"{kwargs['present_v']}_bf16" if kwargs.get("present_v") else ""
+        bf16_outputs = [out_bf16, pk_bf16, pv_bf16]
+
+        self.make_node(
+            "GroupQueryAttention",
+            inputs=bf16_inputs,
+            outputs=bf16_outputs,
+            name=name,
+            domain="com.microsoft",
+            num_heads=self.num_attn_heads,
+            kv_num_heads=self.num_kv_heads,
+            scale=self.attention_attrs["scale"],
+            local_window_size=self.window_size,
+            softcap=self.attention_attrs["softcap"],
+            do_rotary=self.attention_attrs["use_rope_in_attn"],
+            rotary_interleaved=self.rotemb_attrs["interleaved"],
+        )
+
+        self.make_node(
+            "Cast",
+            inputs=[out_bf16],
+            outputs=[f"{name}/output_0"],
+            name=f"{name}/Cast_output_to_fp32",
+            to=TensorProto.FLOAT,
+        )
+
+        if pk_bf16:
+            self.make_node(
+                "Cast",
+                inputs=[pk_bf16],
+                outputs=[kwargs["present_k"]],
+                name=f"{name}/Cast_present_k_to_fp32",
+                to=TensorProto.FLOAT,
+            )
+
+        if pv_bf16:
+            self.make_node(
+                "Cast",
+                inputs=[pv_bf16],
+                outputs=[kwargs["present_v"]],
+                name=f"{name}/Cast_present_v_to_fp32",
+                to=TensorProto.FLOAT,
+            )
+
+        self.make_value_info(
+            f"{name}/output_0",
+            TensorProto.FLOAT,
+            shape=['batch_size', 'sequence_length', self.head_size * self.num_attn_heads]
+        )
+
     def make_sparse_attention(self, name, **kwargs):
         inputs = [
             kwargs["q_path"], kwargs["k_path"], kwargs["v_path"],
@@ -3552,4 +3639,4 @@ def get_args():
 if __name__ == '__main__':
     args = get_args()
     extra_options = parse_extra_options(args.extra_options)
-    create_model(args.model_name, args.input, args.output, args.precision, args.execution_provider, args.cache_dir, **extra_options)
+    create_model(args.model_name, args.input, args.output, args.precision, args.execution_provider, args.cache_dir, **extra_options)
\ No newline at end of file

From 1e541b6be25ff25e4701b01a82e3f19b9075c741 Mon Sep 17 00:00:00 2001
From: Nenad Banfic <nebanfic@microsoft.com>
Date: Tue, 29 Apr 2025 19:49:11 +0000
Subject: [PATCH 13/26] Fast Gelu

---
 src/python/py/models/builder.py | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index a93cc2bcde..1effaf4401 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -2207,6 +2207,24 @@ def make_gelu(self, layer_id, root_input, activation):
 
         return gelu_name
 
+    def make_gelu_bf16(self, layer_id, root_input, activation):
+        gelu_name = f"/model/layers.{layer_id}/mlp/act_fn/{activation}"
+        cast_to_bf16 = f"{gelu_name}/cast_to_bf16"
+        gelu_in = f"{cast_to_bf16}/output"
+        gelu_out_bf16 = f"{gelu_name}/gelu_bf16"
+        cast_to_fp32 = f"{gelu_name}/cast_to_fp32"
+        final_output =  f"{gelu_name}/output_0"
+
+        self.make_node("Cast", inputs=[root_input], outputs=[gelu_in], name=cast_to_bf16, to=TensorProto.BFLOAT16)
+        self.make_node(activation, inputs=[gelu_in], outputs=[gelu_out_bf16], name=gelu_name, domain="com.microsoft")
+        
+        self.make_value_info(gelu_out_bf16, TensorProto.BFLOAT16, shape=['batch_size', 'sequence_length', self.intermediate_size])
+
+        self.make_node("Cast",inputs=[gelu_out_bf16], outputs=[final_output], name=cast_to_fp32, to=TensorProto.FLOAT)
+        self.make_value_info(final_output, self.io_dtype, shape=['batch_size', 'sequence_length', self.intermediate_size])
+
+        return gelu_name
+
     def make_relu(self, layer_id, root_input, activation):
         relu_name = f"/model/layers.{layer_id}/mlp/act_fn/{activation}"
         output = f"{relu_name}/output_0"
@@ -2227,7 +2245,7 @@ def make_activation(self, layer_id, root_input):
         if self.activation in {"silu", "swish", "swiglu"}:
             output_name = self.make_activation_with_mul(layer_id, root_input, activation="Sigmoid", domain=None)
         elif self.activation in {"gelu_new", "gelu_fast", "gelu_pytorch_tanh"}:
-            output_name = self.make_gelu(layer_id, root_input, activation="FastGelu")
+            output_name = self.make_gelu_bf16(layer_id, root_input, activation="FastGelu")
         elif self.activation in {"gelu"}:
             output_name = self.make_gelu(layer_id, root_input, activation="Gelu")
         elif self.activation in {"gegelu", "geglu"}:

From be4e533e225523bad7ed1a11e3ce037821a81edd Mon Sep 17 00:00:00 2001
From: Nenad Banfic <46795300+nenad1002@users.noreply.github.com>
Date: Fri, 2 May 2025 11:02:25 -0700
Subject: [PATCH 14/26] Add Gemma3 bf16 related changes (#1441)

The Gemma3 model must compute layer norms and residual connections in
FP32 to ensure accurate logits. Additionally, all logits are now
explicitly cast to FP32. This PR removes the incorrect code paths and
refactors the builder script so that it cleanly supports both Gemma and
non-Gemma models.
---
 src/python/py/models/builder.py | 249 +++++++++++++++-----------------
 1 file changed, 117 insertions(+), 132 deletions(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index 1effaf4401..30ef8741ee 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -100,9 +100,12 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
 
         # Map output names to their types and shapes
         self.output_names = ["logits"]
+
+        logits_type = TensorProto.FLOAT if "gemma-3" in self.model_name_or_path else self.io_dtype
+
         self.output_types = {
             "hidden_states": self.io_dtype,                                                                      # For standard models where you want to remove the language modeling head from the model (note that `hidden_states` is written this way to match Hugging Face format)
-            "logits": self.io_dtype,                                                                             # For standard models
+            "logits": logits_type,                                                                               # For standard models
             "present.key": self.io_dtype,                                                                        # For standard models (note that `present.key` is written this way to match Hugging Face format)
             "present.value": self.io_dtype,                                                                      # For standard models (note that `present.value` is written this way to match Hugging Face format)
         }
@@ -799,24 +802,20 @@ def make_matmul_op(self, matmul, basename, root_input, **kwargs):
 
     def make_matmul_float(self, matmul, name, root_input, **kwargs):
         weight = name[1:].replace("/", ".") + ".weight"
-        self.make_external_tensor(matmul.weight.detach().T.to(torch.bfloat16).contiguous(), weight)
+        self.make_external_tensor(matmul.weight.detach().T.to( self.to_torch_dtype[self.io_dtype]).contiguous(), weight)
 
         last_dim = matmul.weight.shape[0]
-        first_dim = matmul.weight.shape[1]
         output = "logits" if kwargs.get("logits", False) else f"{name}/output_0"
 
-        cast_input = f"{root_input}/Cast"
-        cast_output = f"{output}/Cast"
-
-        if kwargs.get("first_cast", True):
-            self.make_node("Cast", inputs=[root_input], outputs=[cast_input], name=f"{name}/Cast_Input", to=TensorProto.BFLOAT16)
-            self.make_value_info(cast_input, TensorProto.BFLOAT16, shape=["batch_size", "sequence_length", first_dim])
-
-        self.make_node("MatMul", inputs=[cast_input, weight], outputs=[cast_output], name=name)
-        self.make_value_info(output, self.io_dtype, shape=['batch_size', 'sequence_length', last_dim])
-
-        self.make_node("Cast", inputs=[cast_output], outputs=[output], name=f"{name}/Cast_Output", to=self.io_dtype)
-        self.make_value_info(cast_output, TensorProto.BFLOAT16, shape=["batch_size", "sequence_length", last_dim])
+        if output != "logits" or "gemma-3" not in self.model_name_or_path:
+            self.make_node("MatMul", inputs=[root_input, weight], outputs=[output], name=name)
+            self.make_value_info(output, self.io_dtype, shape=['batch_size', 'sequence_length', last_dim])
+        else:
+            cast_name = name + "/Cast"
+            cast_name_input = name + "/Cast"
+            self.make_node("MatMul", inputs=[root_input, weight], outputs=[cast_name_input], name=name)
+            self.make_node("Cast", inputs=[cast_name_input], outputs=[output], name=cast_name, to=TensorProto.FLOAT)
+            self.make_value_info(output, TensorProto.FLOAT, shape=['batch_size', 'sequence_length', last_dim])
 
         return name
 
@@ -1068,7 +1067,7 @@ def make_embedding(self, embedding):
         self.layernorm_attrs["root_input"] = layernorm_attrs_value
         self.layernorm_attrs["skip_input"] = layernorm_attrs_value
 
-    def make_layernorm(self, layer_id, layernorm, skip, simple, location):
+    def make_layernorm_default(self, layer_id, layernorm, skip, simple, location):
         root_input = self.layernorm_attrs["root_input"]
         skip_input = self.layernorm_attrs["skip_input"]
 
@@ -1107,6 +1106,74 @@ def make_layernorm(self, layer_id, layernorm, skip, simple, location):
             # Assign output 3 of current SkipLayerNorm as root input to next SkipLayerNorm
             self.layernorm_attrs["root_input"] = output_3
 
+    def make_layernorm_with_casting_to_fp32(self, layer_id, layernorm, skip, simple, location):
+        root_input = self.layernorm_attrs["root_input"]
+        skip_input = self.layernorm_attrs["skip_input"]
+
+        skip_path = "skip" if skip else "noskip"
+
+        casted_root = f"{root_input}.{layer_id}.{skip_path}root/Cast"
+        self.make_node("Cast", inputs=[root_input], outputs=[casted_root], name=f"/model/layers.{layer_id}/{location}_layernorm/CastInput0", to=TensorProto.FLOAT)
+        root_input = casted_root
+
+        self.make_value_info(casted_root, TensorProto.FLOAT, shape=["batch_size", "sequence_length", self.hidden_size],)
+
+        if skip:
+            casted_skip = f"{skip_input}.{layer_id}.{skip_path}skip/Cast"
+            self.make_node("Cast", inputs=[skip_input], outputs=[casted_skip], name=f"/model/layers.{layer_id}/{location}_layernorm/CastInput3", to=TensorProto.FLOAT)
+            skip_input = casted_skip
+
+            self.make_value_info(casted_skip, TensorProto.FLOAT, shape=["batch_size", "sequence_length", self.hidden_size],)
+
+        weight = f"model.layers.{layer_id}.{location}_layernorm.weight"
+        self.make_external_tensor(layernorm.weight.detach().to(self.to_torch_dtype[TensorProto.FLOAT]).contiguous() + self.layernorm_attrs["add_offset"], weight)
+        bias = f"model.layers.{layer_id}.{location}_layernorm.bias"
+        if not simple:
+            self.make_external_tensor(layernorm.bias.detach().to(self.to_torch_dtype[TensorProto.FLOAT]).contiguous(), bias)
+
+        inputs = [root_input, skip_input, weight] if skip else [root_input, weight]
+        if not simple:
+            inputs.append(bias)
+
+        name = f"/model/layers.{layer_id}/{location}_layernorm/{'Skip' if skip else ''}LayerNorm"
+        op_type = f"{'Skip' if skip else ''}{'Simplified' if simple else ''}LayerNormalization"
+        kwargs = {"epsilon": self.layernorm_attrs["epsilon"]}
+        if not skip:
+            kwargs.update({"axis": -1, "stash_type": 1})
+
+        output_0 = f"/model/layers.{layer_id}/{location}_layernorm/output_0"
+        output_3 = f"/model/layers.{layer_id}/{location}_layernorm/output_3"
+        raw_out0 = f"{output_0}/Cast"
+        raw_out3 = f"{output_3}/Cast"
+        if self.layernorm_attrs["last_layernorm"] and (self.include_hidden_states or self.exclude_lm_head):
+            raw_out0 = "hidden_states"
+        outputs = [raw_out0, "", "", raw_out3] if skip and not self.layernorm_attrs["last_layernorm"] else [raw_out0]
+
+        self.make_node(op_type, inputs=inputs, outputs=outputs, name=name, domain=("com.microsoft" if skip else None), **kwargs)
+
+        self.make_value_info(raw_out0, TensorProto.FLOAT, shape=['batch_size', 'sequence_length', self.hidden_size])
+        if skip and not self.layernorm_attrs["last_layernorm"]:
+            self.make_value_info(raw_out3, TensorProto.FLOAT, shape=['batch_size', 'sequence_length', self.hidden_size])
+
+        if not(skip and not self.layernorm_attrs["last_layernorm"]):
+            raw_out3 = None
+
+        self.make_node("Cast", inputs=[raw_out0], outputs=[output_0], name=f"{name}/CastOutput0", to=self.io_dtype,)
+        self.make_value_info(output_0, self.io_dtype, shape=["batch_size", "sequence_length", self.hidden_size],)
+        self.layernorm_attrs["output_0"] = output_0
+
+        if raw_out3 is not None:
+            self.make_node("Cast", inputs=[raw_out3], outputs=[output_3], name=f"{name}/CastOutput3", to=self.io_dtype,)
+            self.make_value_info(output_3, self.io_dtype, shape=["batch_size", "sequence_length", self.hidden_size],)
+            self.layernorm_attrs["output_3"] = output_3
+            self.layernorm_attrs["root_input"] = output_3
+
+    def make_layernorm(self, layer_id, layernorm, skip, simple, location):
+        if "gemma-3" in self.model_name_or_path and self.io_dtype == TensorProto.BFLOAT16:
+            self.make_layernorm_with_casting_to_fp32(layer_id, layernorm, skip, simple, location)
+        else:
+            self.make_layernorm_default(layer_id, layernorm, skip, simple, location)
+
     def make_mscale_su(self, mscale):
         if mscale <= 1.0:
             return 1.0
@@ -1171,9 +1238,9 @@ def make_rotary_embedding_caches(self, **kwargs):
             # Slice cos/sin caches from (M, H) to (M, H/2)
             hidden_dim = cos_cache.shape[-1]
             cos_cache = cos_cache.squeeze()[:, : (hidden_dim // 2)].detach()
-            cos_cache = cos_cache.to(self.to_torch_dtype[TensorProto.BFLOAT16])
+            cos_cache = cos_cache.to(self.to_torch_dtype[self.io_dtype])
             sin_cache = sin_cache.squeeze()[:, : (hidden_dim // 2)].detach()
-            sin_cache = sin_cache.to(self.to_torch_dtype[TensorProto.BFLOAT16])
+            sin_cache = sin_cache.to(self.to_torch_dtype[self.io_dtype])
 
             # Slice cos/sin caches from (M, H/2) to (M, R/2) if partial rotary embeddings are used
             if self.rotemb_attrs["partial_rotary_factor"] != 1.0:
@@ -1196,23 +1263,15 @@ def make_rotary_embedding(self, name, root_input, **kwargs):
         cos_cache_name, sin_cache_name = self.make_rotary_embedding_caches()
         num_heads = self.num_kv_heads if "k_rotary" in name else self.num_attn_heads
 
-        casted_root = f"{name}/root_input_bf16"
-        self.make_node(
-            "Cast",
-            inputs=[root_input],
-            outputs=[casted_root],
-            name=f"{name}/Cast_to_bf16",
-            to=TensorProto.BFLOAT16,
-        )
+        inputs = [root_input, kwargs.pop("position_ids"), cos_cache_name, sin_cache_name]
 
-        inputs = [casted_root, kwargs.pop("position_ids"), cos_cache_name, sin_cache_name]
         output = f"{name}/output_0"
         self.make_node(
             "RotaryEmbedding", inputs=inputs, outputs=[output], name=name, domain="com.microsoft",
             interleaved=self.rotemb_attrs["interleaved"], num_heads=(0 if self.rotemb_attrs["partial_rotary_factor"] == 1.0 else num_heads),  # default is 0 in RotaryEmbedding kernel
             rotary_embedding_dim=self.rotemb_attrs["rotary_embedding_dim"],
         )
-        self.make_value_info(output, TensorProto.BFLOAT16, shape=['batch_size', 'sequence_length', self.head_size * num_heads])
+        self.make_value_info(output, self.io_dtype, shape=['batch_size', 'sequence_length', self.head_size * num_heads])
 
     def make_rotary_embedding_multi_cache(self, **kwargs):
         cos_cache_name = kwargs.get("cos_cache_name", "cos_cache")
@@ -1320,12 +1379,23 @@ def make_qk_norm(self, layer_id, attention):
         q_reshape_1_output = f"{q_reshape_1_name}/output_0"
         self.make_reshape(q_reshape_1_name, q_reshape_1_inputs, dtype=self.io_dtype, shape=['batch_size', 'sequence_length * num_attention_heads', self.head_size])
 
-        # Make Q LayerNorm
         q_layernorm_name = f"/model/layers.{layer_id}/attn/q_norm/SimplifiedLayerNormalization"
         q_weight_name = f"model.layers.{layer_id}.attn.q_norm.layernorm.weight"
         q_layernorm_output = f"{q_layernorm_name}/output_0"
-        self.make_external_tensor((attention.q_norm.weight.detach().to(self.to_torch_dtype[self.io_dtype]) + self.layernorm_attrs["add_offset"]).contiguous(), q_weight_name)
-        self.make_node("SimplifiedLayerNormalization", inputs=[q_reshape_1_output, q_weight_name], outputs=[q_layernorm_output], name=q_layernorm_name, **layernorm_kwargs)
+
+        if "gemma-3" in self.model_name_or_path and self.io_dtype == TensorProto.BFLOAT16:
+            q_layernorm_cast_in = f"{q_layernorm_name}/Cast"
+            self.make_node("Cast", inputs=[q_reshape_1_output], outputs=[q_layernorm_cast_in], name=f"{q_layernorm_name}/CastIn", to=TensorProto.FLOAT)
+
+            q_layernorm_weight_fp32 = (attention.q_norm.weight.detach().to(self.to_torch_dtype[TensorProto.FLOAT]) + self.layernorm_attrs["add_offset"]).contiguous()
+            self.make_external_tensor(q_layernorm_weight_fp32, q_weight_name)
+
+            self.make_node("SimplifiedLayerNormalization", inputs=[q_layernorm_cast_in, q_weight_name], outputs=[f"{q_layernorm_name}/output/Cast"], name=q_layernorm_name, **layernorm_kwargs)
+            self.make_node("Cast", inputs=[f"{q_layernorm_name}/output/Cast"], outputs=[q_layernorm_output], name=f"{q_layernorm_name}/CastOut", to=self.io_dtype)
+        else:
+            self.make_external_tensor((attention.q_norm.weight.detach().to(self.to_torch_dtype[self.io_dtype]) + self.layernorm_attrs["add_offset"]).contiguous(), q_weight_name)
+            self.make_node("SimplifiedLayerNormalization", inputs=[q_reshape_1_output, q_weight_name], outputs=[q_layernorm_output], name=q_layernorm_name, **layernorm_kwargs)
+       
         self.make_value_info(q_layernorm_output, dtype=self.io_dtype, shape=['batch_size', 'sequence_length * num_attention_heads', self.head_size])
 
         # Reshape Q path after LayerNorm from Bx(SxN)xH to BxSxD
@@ -1339,12 +1409,23 @@ def make_qk_norm(self, layer_id, attention):
         k_reshape_1_output = f"{k_reshape_1_name}/output_0"
         self.make_reshape(k_reshape_1_name, k_reshape_1_inputs, dtype=self.io_dtype, shape=['batch_size', 'sequence_length * num_key_value_heads', self.head_size])
 
-        # Make K LayerNorm
         k_layernorm_name = f"/model/layers.{layer_id}/attn/k_norm/SimplifiedLayerNormalization"
         k_weight_name = f"model.layers.{layer_id}.attn.k_norm.layernorm.weight"
         k_layernorm_output = f"{k_layernorm_name}/output_0"
-        self.make_external_tensor((attention.k_norm.weight.detach().to(self.to_torch_dtype[self.io_dtype]) + self.layernorm_attrs["add_offset"]).contiguous(), k_weight_name)
-        self.make_node("SimplifiedLayerNormalization", inputs=[k_reshape_1_output, k_weight_name], outputs=[k_layernorm_output], name=k_layernorm_name, **layernorm_kwargs)
+
+        if "gemma-3" in self.model_name_or_path and self.io_dtype == TensorProto.BFLOAT16:
+            k_layernorm_cast_in = f"{k_layernorm_name}/Cast"
+            self.make_node("Cast", inputs=[k_reshape_1_output], outputs=[k_layernorm_cast_in], name=f"{k_layernorm_name}/CastIn", to=TensorProto.FLOAT)
+
+            k_layernorm_weight_fp32 = (attention.k_norm.weight.detach().to(self.to_torch_dtype[TensorProto.FLOAT]) + self.layernorm_attrs["add_offset"]).contiguous()
+            self.make_external_tensor(k_layernorm_weight_fp32, k_weight_name)
+
+            self.make_node("SimplifiedLayerNormalization", inputs=[k_layernorm_cast_in, k_weight_name], outputs=[f"{k_layernorm_name}/output/Cast"], name=k_layernorm_name, **layernorm_kwargs)
+            self.make_node("Cast", inputs=[f"{k_layernorm_name}/output/Cast"], outputs=[k_layernorm_output], name=f"{k_layernorm_name}/CastOut", to=self.io_dtype)
+        else:
+            self.make_external_tensor((attention.k_norm.weight.detach().to(self.to_torch_dtype[self.io_dtype]) + self.layernorm_attrs["add_offset"]).contiguous(), k_weight_name)
+            self.make_node("SimplifiedLayerNormalization", inputs=[k_reshape_1_output, k_weight_name], outputs=[k_layernorm_output], name=k_layernorm_name, **layernorm_kwargs)                
+            
         self.make_value_info(k_layernorm_output, dtype=self.io_dtype, shape=['batch_size', 'sequence_length * num_key_value_heads', self.head_size])
 
         # Reshape K path after LayerNorm from Bx(SxN)xH to BxSxD
@@ -1540,7 +1621,7 @@ def make_attention_op(self, name, **kwargs):
         if op_type == "MultiHeadAttention":
             self.make_multi_head_attention(name, add_qk=f"{self.mask_attrs['mask_name']}/output_0", **kwargs)
         elif op_type == "GroupQueryAttention":
-            self.make_group_query_attention_with_bf16(name, seqlens_k=f"{self.mask_attrs['seqlens_k']}/output_0", total_seq_len=f"{self.mask_attrs['total_seq_len']}/output_0", **kwargs)
+            self.make_group_query_attention(name, seqlens_k=f"{self.mask_attrs['seqlens_k']}/output_0", total_seq_len=f"{self.mask_attrs['total_seq_len']}/output_0", **kwargs)
         elif op_type == "SparseAttention":
             self.make_sparse_attention(name, block_row_indices=self.mask_attrs['block_row_indices'], block_col_indices=self.mask_attrs['block_col_indices'], key_total_seq_lens=f"{self.mask_attrs['key_total_seq_lens']}/output_0", total_seq_len=f"{self.mask_attrs['total_seq_len']}/output_0", **kwargs)
         else:
@@ -1576,84 +1657,6 @@ def make_group_query_attention(self, name, **kwargs):
         )
         self.make_value_info(output, self.io_dtype, shape=['batch_size', 'sequence_length', self.head_size * self.num_attn_heads])
 
-    def make_group_query_attention_with_bf16(self, name, **kwargs):
-        raw_inputs = [
-            kwargs["v_path"],
-            kwargs.get("past_k", ""), kwargs.get("past_v", ""),
-            kwargs.get("cos_cache", ""), kwargs.get("sin_cache", ""),
-        ]
-
-        int64_inputs = [kwargs.get("seqlens_k", ""), kwargs.get("total_seq_len", "")]
-
-        bf16_inputs = [kwargs["q_path"], kwargs["k_path"]]
-        for inp in raw_inputs:
-            if inp == "": 
-                bf16_inputs.append("") 
-                continue
-            bf16_inp = f"{inp}_bf16"
-            self.make_node(
-                "Cast",
-                inputs=[inp],
-                outputs=[bf16_inp],
-                name=f"{name}/cast_{inp}_to_bf16",
-                to=TensorProto.BFLOAT16
-            )
-            bf16_inputs.append(bf16_inp)
-
-        bf16_inputs[5:5] = int64_inputs
-
-        out_bf16 = f"{name}/output_0_bf16"
-        pk_bf16  = f"{kwargs['present_k']}_bf16" if kwargs.get("present_k") else ""
-        pv_bf16  = f"{kwargs['present_v']}_bf16" if kwargs.get("present_v") else ""
-        bf16_outputs = [out_bf16, pk_bf16, pv_bf16]
-
-        self.make_node(
-            "GroupQueryAttention",
-            inputs=bf16_inputs,
-            outputs=bf16_outputs,
-            name=name,
-            domain="com.microsoft",
-            num_heads=self.num_attn_heads,
-            kv_num_heads=self.num_kv_heads,
-            scale=self.attention_attrs["scale"],
-            local_window_size=self.window_size,
-            softcap=self.attention_attrs["softcap"],
-            do_rotary=self.attention_attrs["use_rope_in_attn"],
-            rotary_interleaved=self.rotemb_attrs["interleaved"],
-        )
-
-        self.make_node(
-            "Cast",
-            inputs=[out_bf16],
-            outputs=[f"{name}/output_0"],
-            name=f"{name}/Cast_output_to_fp32",
-            to=TensorProto.FLOAT,
-        )
-
-        if pk_bf16:
-            self.make_node(
-                "Cast",
-                inputs=[pk_bf16],
-                outputs=[kwargs["present_k"]],
-                name=f"{name}/Cast_present_k_to_fp32",
-                to=TensorProto.FLOAT,
-            )
-
-        if pv_bf16:
-            self.make_node(
-                "Cast",
-                inputs=[pv_bf16],
-                outputs=[kwargs["present_v"]],
-                name=f"{name}/Cast_present_v_to_fp32",
-                to=TensorProto.FLOAT,
-            )
-
-        self.make_value_info(
-            f"{name}/output_0",
-            TensorProto.FLOAT,
-            shape=['batch_size', 'sequence_length', self.head_size * self.num_attn_heads]
-        )
-
     def make_sparse_attention(self, name, **kwargs):
         inputs = [
             kwargs["q_path"], kwargs["k_path"], kwargs["v_path"],
@@ -2207,24 +2210,6 @@ def make_gelu(self, layer_id, root_input, activation):
 
         return gelu_name
 
-    def make_gelu_bf16(self, layer_id, root_input, activation):
-        gelu_name = f"/model/layers.{layer_id}/mlp/act_fn/{activation}"
-        cast_to_bf16 = f"{gelu_name}/cast_to_bf16"
-        gelu_in = f"{cast_to_bf16}/output"
-        gelu_out_bf16 = f"{gelu_name}/gelu_bf16"
-        cast_to_fp32 = f"{gelu_name}/cast_to_fp32"
-        final_output =  f"{gelu_name}/output_0"
-
-        self.make_node("Cast", inputs=[root_input], outputs=[gelu_in], name=cast_to_bf16, to=TensorProto.BFLOAT16)
-        self.make_node(activation, inputs=[gelu_in], outputs=[gelu_out_bf16], name=gelu_name, domain="com.microsoft")
-        
-        self.make_value_info(gelu_out_bf16, TensorProto.BFLOAT16, shape=['batch_size', 'sequence_length', self.intermediate_size])
-
-        self.make_node("Cast",inputs=[gelu_out_bf16], outputs=[final_output], name=cast_to_fp32, to=TensorProto.FLOAT)
-        self.make_value_info(final_output, self.io_dtype, shape=['batch_size', 'sequence_length', self.intermediate_size])
-
-        return gelu_name
-
     def make_relu(self, layer_id, root_input, activation):
         relu_name = f"/model/layers.{layer_id}/mlp/act_fn/{activation}"
         output = f"{relu_name}/output_0"
@@ -2245,7 +2230,7 @@ def make_activation(self, layer_id, root_input):
         if self.activation in {"silu", "swish", "swiglu"}:
             output_name = self.make_activation_with_mul(layer_id, root_input, activation="Sigmoid", domain=None)
         elif self.activation in {"gelu_new", "gelu_fast", "gelu_pytorch_tanh"}:
-            output_name = self.make_gelu_bf16(layer_id, root_input, activation="FastGelu")
+            output_name = self.make_gelu(layer_id, root_input, activation="FastGelu")
         elif self.activation in {"gelu"}:
             output_name = self.make_gelu(layer_id, root_input, activation="Gelu")
         elif self.activation in {"gegelu", "geglu"}:

From 1ca152d843e1ae8b5208230db273b2ee0d16ba97 Mon Sep 17 00:00:00 2001
From: Nenad Banfic <nebanfic@microsoft.com>
Date: Wed, 7 May 2025 16:49:22 +0000
Subject: [PATCH 15/26] Add layernorms

---
 src/python/py/models/builder.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index 30ef8741ee..1bf4a36c32 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -1169,7 +1169,7 @@ def make_layernorm_with_casting_to_fp32(self, layer_id, layernorm, skip, simple,
             self.layernorm_attrs["root_input"] = output_3
 
     def make_layernorm(self, layer_id, layernorm, skip, simple, location):
-        if "gemma-3" in self.model_name_or_path and self.io_dtype == TensorProto.BFLOAT16:
+        if self.layernorm_attrs["use_fp32_layernorms"]:
             self.make_layernorm_with_casting_to_fp32(layer_id, layernorm, skip, simple, location)
         else:
             self.make_layernorm_default(layer_id, layernorm, skip, simple, location)
@@ -2325,6 +2325,7 @@ def make_model(self, input_path):
                 from onnxruntime_genai.models.gguf_model import GGUFModel
             model = GGUFModel.from_pretrained(self.model_type, input_path, self.head_size, self.hidden_size, self.intermediate_size, self.num_attn_heads, self.num_kv_heads, self.vocab_size)
             self.layernorm_attrs["add_offset"] = 0  # add offset already done for GGUF models
+            self.layernorm_attrs["use_fp32_layernorms"] = 0
         elif self.quant_type is not None:
             # Load quantized PyTorch model
             try:
@@ -3334,6 +3335,7 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
         self.is_local = lambda layer_id: bool((layer_id + 1) % config.sliding_window_pattern)
         self.rope_local_theta = config.rope_local_base_freq
         self.make_rotary_embedding_multi_cache()
+        self.layernorm_attrs["use_fp32_layernorms"] = 1 if self.io_dtype == TensorProto.BFLOAT16 else 0        
 
     def make_attention_init(self):
         self.attention_attrs["q_norm"] = True

From b1e437803d6cd44319e2a28f77705412c34dc609 Mon Sep 17 00:00:00 2001
From: Kunal Vaishnavi <kvaishnavi@microsoft.com>
Date: Wed, 7 May 2025 20:30:40 +0000
Subject: [PATCH 16/26] Fix how final norm is accessed

---
 src/python/py/models/builder.py | 13 ++++++++++---
 1 file changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index b56766359f..88bbdea395 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -2365,7 +2365,15 @@ def make_model(self, input_path):
         del model
 
     def has_final_norm(self, module, orig_model):
-        if hasattr(orig_model, "base_model") and hasattr(orig_model.base_model, "model"):
+        # Find where the language model is stored to check attributes. Some classes
+        # store the language model in a different attribute than `model.model`.
+        if hasattr(orig_model, "language_model"):
+            # Model is multimodal
+            # Note: This case is checked first because the `language_model` attribute and the `base_model` attribute
+            # exist for both multimodal models and PEFT models. However they represent different classes and their attributes
+            # differ.
+            model = orig_model.language_model
+        elif hasattr(orig_model, "base_model") and hasattr(orig_model.base_model, "model"):
             # Model is from PEFT
             model = orig_model.base_model.model
         else:
@@ -2375,12 +2383,11 @@ def has_final_norm(self, module, orig_model):
         hf_norm = hasattr(model, "model") and hasattr(model.model, "norm") and module == model.model.norm
         hf_final_layernorm = hasattr(model, "model") and hasattr(model.model, "final_layernorm") and module == model.model.final_layernorm
         hf_transformer_final_layernorm = hasattr(model, "transformer") and hasattr(model.transformer, "encoder") and hasattr(model.transformer.encoder, "final_layernorm") and module == model.transformer.encoder.final_layernorm
-        hf_multimodal_final_layernorm = hasattr(model, "language_model") and hasattr(model.language_model, "model") and hasattr(model.language_model.model, "norm") and module == model.language_model.model.norm
 
         # GGUF names
         gguf_final_norm = hasattr(model, "final_norm") and module == model.final_norm
 
-        hf_names = [hf_norm, hf_final_layernorm, hf_transformer_final_layernorm, hf_multimodal_final_layernorm]
+        hf_names = [hf_norm, hf_final_layernorm, hf_transformer_final_layernorm]
         gguf_names = [gguf_final_norm]
         return any(hf_names + gguf_names)
 

From 19654b2269d041f03530e5b11650d348b5b953b1 Mon Sep 17 00:00:00 2001
From: Nenad Banfic <nebanfic@microsoft.com>
Date: Wed, 7 May 2025 22:45:05 +0000
Subject: [PATCH 17/26] Add guard methods

---
 src/python/py/models/builder.py | 58 ++++++++++++++++++++-------------
 1 file changed, 36 insertions(+), 22 deletions(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index 88bbdea395..7e57fcace3 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -101,20 +101,9 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
         # Map output names to their types and shapes
         self.output_names = ["logits"]
 
-        logits_type = TensorProto.FLOAT if "gemma-3" in self.model_name_or_path else self.io_dtype
+        self.outputs_attrs = {"logits_type": self.io_dtype}
 
-        self.output_types = {
-            "hidden_states": self.io_dtype,                                                                      # For standard models where you want to remove the language modeling head from the model (note that `hidden_states` is written this way to match Hugging Face format)
-            "logits": logits_type,                                                                               # For standard models
-            "present.key": self.io_dtype,                                                                        # For standard models (note that `present.key` is written this way to match Hugging Face format)
-            "present.value": self.io_dtype,                                                                      # For standard models (note that `present.value` is written this way to match Hugging Face format)
-        }
-        self.output_shapes = {
-            "hidden_states": ["batch_size", "sequence_length", self.hidden_size],                                # For standard models where you want to remove the language modeling head from the model (note that `hidden_states` is written this way to match Hugging Face format)
-            "logits": ["batch_size", "sequence_length", self.vocab_size],                                        # For standard models
-            "present.key": ["batch_size", self.num_kv_heads, "total_sequence_length", self.head_size],           # For standard models (note that `present.key` is written this way to match Hugging Face format)
-            "present.value": ["batch_size", self.num_kv_heads, "total_sequence_length", self.head_size],         # For standard models (note that `present.value` is written this way to match Hugging Face format)
-        }
+        self.make_outputs_init()
         self.exclude_lm_head = extra_options.get("exclude_lm_head", False)
         self.include_hidden_states = extra_options.get("include_hidden_states", False)
         if self.exclude_lm_head:
@@ -320,6 +309,21 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
             self.quant_attrs["config"] = config.quantization_config
             self.quant_attrs["use_g_idx"] = config.quantization_config["desc_act"] if "desc_act" in config.quantization_config else False
 
+
+    def make_outputs_init(self):
+        self.output_types = {
+            "hidden_states": self.io_dtype,                                                                      # For standard models where you want to remove the language modeling head from the model (note that `hidden_states` is written this way to match Hugging Face format)
+            "logits": self.outputs_attrs["logits_type"],                                                                               # For standard models
+            "present.key": self.io_dtype,                                                                        # For standard models (note that `present.key` is written this way to match Hugging Face format)
+            "present.value": self.io_dtype,                                                                      # For standard models (note that `present.value` is written this way to match Hugging Face format)
+        }
+        self.output_shapes = {
+            "hidden_states": ["batch_size", "sequence_length", self.hidden_size],                                # For standard models where you want to remove the language modeling head from the model (note that `hidden_states` is written this way to match Hugging Face format)
+            "logits": ["batch_size", "sequence_length", self.vocab_size],                                        # For standard models
+            "present.key": ["batch_size", self.num_kv_heads, "total_sequence_length", self.head_size],           # For standard models (note that `present.key` is written this way to match Hugging Face format)
+            "present.value": ["batch_size", self.num_kv_heads, "total_sequence_length", self.head_size],         # For standard models (note that `present.value` is written this way to match Hugging Face format)
+        }
+
     def make_attention_init(self):
         valid_gqa_configurations = [
             ("cpu", TensorProto.FLOAT),
@@ -820,18 +824,14 @@ def make_matmul_float(self, matmul, name, root_input, **kwargs):
         last_dim = matmul.weight.shape[0]
         output = "logits" if kwargs.get("logits", False) else f"{name}/output_0"
 
-        if output != "logits" or "gemma-3" not in self.model_name_or_path:
-            self.make_node("MatMul", inputs=[root_input, weight], outputs=[output], name=name)
-            self.make_value_info(output, self.io_dtype, shape=['batch_size', 'sequence_length', last_dim])
-        else:
-            cast_name = name + "/Cast"
-            cast_name_input = name + "/Cast"
-            self.make_node("MatMul", inputs=[root_input, weight], outputs=[cast_name_input], name=name)
-            self.make_node("Cast", inputs=[cast_name_input], outputs=[output], name=cast_name, to=TensorProto.FLOAT)
-            self.make_value_info(output, TensorProto.FLOAT, shape=['batch_size', 'sequence_length', last_dim])
+        self.make_matmul_float_node_and_value_info(name, [root_input, weight], output, last_dim)
 
         return name
 
+    def make_matmul_float_node_and_value_info(self, name, inputs, output, last_dim):
+        self.make_node("MatMul", inputs=inputs, outputs=[output], name=name)
+        self.make_value_info(output, self.io_dtype, shape=['batch_size', 'sequence_length', last_dim])
+
     def make_matmul_int4(self, matmul, basename, root_input, **kwargs):
         if not hasattr(matmul, "qweight"):
             # TODO: quantize weights, then save new MatMul weights for onnx model
@@ -3347,6 +3347,20 @@ def make_attention_init(self):
         self.attention_attrs["k_norm"] = True
         super().make_attention_init()
 
+    def make_outputs_init(self):
+        self.outputs_attrs["logits_type"] = TensorProto.FLOAT
+        super().make_outputs_init()
+
+    def make_matmul_float_node_and_value_info(self, name, inputs, output, last_dim):
+        if output == "logits":
+            cast_name = name + "/Cast"
+            cast_name_input = name + "/Cast"
+            self.make_node("MatMul", inputs=inputs, outputs=[cast_name_input], name=name)
+            self.make_node("Cast", inputs=[cast_name_input], outputs=[output], name=cast_name, to=TensorProto.FLOAT)
+            self.make_value_info(output, TensorProto.FLOAT, shape=['batch_size', 'sequence_length', last_dim])
+        else:
+            super().make_matmul_float_node_and_value_info(name, inputs, output, last_dim)
+
     def make_rotary_embedding_multi_cache(self):
         self.cos_cache_global_name, self.sin_cache_global_name = "cos_cache_global", "sin_cache_global"
         super().make_rotary_embedding_caches(cos_cache_name=self.cos_cache_global_name, sin_cache_name=self.sin_cache_global_name)

From ed9eda04f9efde0e402e94f5c7dcb54029457467 Mon Sep 17 00:00:00 2001
From: Nenad Banfic <nebanfic@microsoft.com>
Date: Thu, 8 May 2025 00:14:41 +0000
Subject: [PATCH 18/26] Refactor layernorm

---
 src/python/py/models/builder.py | 92 +++++++++++++++++----------------
 1 file changed, 48 insertions(+), 44 deletions(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index 7e57fcace3..e90599f0d2 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -100,9 +100,7 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
 
         # Map output names to their types and shapes
         self.output_names = ["logits"]
-
         self.outputs_attrs = {"logits_type": self.io_dtype}
-
         self.make_outputs_init()
         self.exclude_lm_head = extra_options.get("exclude_lm_head", False)
         self.include_hidden_states = extra_options.get("include_hidden_states", False)
@@ -1058,17 +1056,11 @@ def make_layernorm_default(self, layer_id, layernorm, skip, simple, location):
         skip_input = self.layernorm_attrs["skip_input"]
 
         weight = f"model.layers.{layer_id}.{location}_layernorm.weight"
-        self.make_external_tensor(layernorm.weight.detach().cpu().to(self.to_torch_dtype[self.io_dtype]).contiguous() + self.layernorm_attrs["add_offset"], weight)
-        bias = f"model.layers.{layer_id}.{location}_layernorm.bias"
-        if not simple:
-            self.make_external_tensor(layernorm.bias.detach().cpu().to(self.to_torch_dtype[self.io_dtype]).contiguous(), bias)
-
         inputs = [root_input, skip_input, weight] if skip else [root_input, weight]
         if not simple:
             inputs.append(bias)
 
         name = f"/model/layers.{layer_id}/{location}_layernorm/{'Skip' if skip else ''}LayerNorm"
-        op_type = f"{'Skip' if skip else ''}{'Simplified' if simple else ''}LayerNormalization"
         kwargs = {"epsilon": self.layernorm_attrs["epsilon"]}
         if not skip:
             kwargs.update({"axis": -1, "stash_type": 1})
@@ -1078,18 +1070,19 @@ def make_layernorm_default(self, layer_id, layernorm, skip, simple, location):
         if self.layernorm_attrs["last_layernorm"] and (self.include_hidden_states or self.exclude_lm_head):
             output_0 = "hidden_states"
         outputs = [output_0, "", "", output_3] if skip and not self.layernorm_attrs["last_layernorm"] else [output_0]
+        
+        bias_name = f"model.layers.{layer_id}.{location}_layernorm.bias"
+        bias = getattr(layernorm, "bias", None)
 
-        self.make_node(op_type, inputs=inputs, outputs=outputs, name=name, domain=("com.microsoft" if skip else None), **kwargs)
+        self.make_layernorm_node_and_tensors(layernorm.weight, weight, bias, bias_name, inputs, outputs, name, TensorProto.FLOAT, skip, simple, kwargs)
         self.make_value_info(output_0, self.io_dtype, shape=['batch_size', 'sequence_length', self.hidden_size])
         if skip and not self.layernorm_attrs["last_layernorm"]:
             self.make_value_info(output_3, self.io_dtype, shape=['batch_size', 'sequence_length', self.hidden_size])
 
-        # Update LayerNorm attributes
         self.layernorm_attrs["output_0"] = output_0
         if skip and not self.layernorm_attrs["last_layernorm"]:
             self.layernorm_attrs["output_3"] = output_3
 
-            # Assign output 3 of current SkipLayerNorm as root input to next SkipLayerNorm
             self.layernorm_attrs["root_input"] = output_3
 
     def make_layernorm_with_casting_to_fp32(self, layer_id, layernorm, skip, simple, location):
@@ -1112,17 +1105,12 @@ def make_layernorm_with_casting_to_fp32(self, layer_id, layernorm, skip, simple,
             self.make_value_info(casted_skip, TensorProto.FLOAT, shape=["batch_size", "sequence_length", self.hidden_size],)
 
         weight = f"model.layers.{layer_id}.{location}_layernorm.weight"
-        self.make_external_tensor(layernorm.weight.detach().to(self.to_torch_dtype[TensorProto.FLOAT]).contiguous() + self.layernorm_attrs["add_offset"], weight)
-        bias = f"model.layers.{layer_id}.{location}_layernorm.bias"
-        if not simple:
-            self.make_external_tensor(layernorm.bias.detach().to(self.to_torch_dtype[TensorProto.FLOAT]).contiguous(), bias)
-
+       
         inputs = [root_input, skip_input, weight] if skip else [root_input, weight]
         if not simple:
             inputs.append(bias)
 
         name = f"/model/layers.{layer_id}/{location}_layernorm/{'Skip' if skip else ''}LayerNorm"
-        op_type = f"{'Skip' if skip else ''}{'Simplified' if simple else ''}LayerNormalization"
         kwargs = {"epsilon": self.layernorm_attrs["epsilon"]}
         if not skip:
             kwargs.update({"axis": -1, "stash_type": 1})
@@ -1135,7 +1123,9 @@ def make_layernorm_with_casting_to_fp32(self, layer_id, layernorm, skip, simple,
             raw_out0 = "hidden_states"
         outputs = [raw_out0, "", "", raw_out3] if skip and not self.layernorm_attrs["last_layernorm"] else [raw_out0]
 
-        self.make_node(op_type, inputs=inputs, outputs=outputs, name=name, domain=("com.microsoft" if skip else None), **kwargs)
+        bias_name = f"model.layers.{layer_id}.{location}_layernorm.bias"
+        bias = getattr(layernorm, "bias", None)
+        self.make_layernorm_node_and_tensors(layernorm.weight, weight, bias, bias_name, inputs, outputs, name, TensorProto.FLOAT, skip, simple, kwargs)
 
         self.make_value_info(raw_out0, TensorProto.FLOAT, shape=['batch_size', 'sequence_length', self.hidden_size])
         if skip and not self.layernorm_attrs["last_layernorm"]:
@@ -1345,6 +1335,44 @@ def make_rotary_embedding_multi_cache(self, **kwargs):
         self.make_value_info(cos_cache_name, self.io_dtype, shape=["max_sequence_length", "head_dim / 2"])
         self.make_value_info(sin_cache_name, self.io_dtype, shape=["max_sequence_length", "head_dim / 2"])
 
+    def make_layernorm_node_and_tensors(self, weight, weight_name, bias, bias_name, inputs, outputs, name, tensor_type, skip, simple, layernorm_kwargs):
+        op_type = f"{'Skip' if skip else ''}{'Simplified' if simple else ''}LayerNormalization"
+        self.make_external_tensor((weight.detach().cpu().to(self.to_torch_dtype[tensor_type]) + self.layernorm_attrs["add_offset"]).contiguous(), weight_name)
+        
+        if not simple:
+            self.make_external_tensor(bias.detach().to(self.to_torch_dtype[TensorProto.FLOAT]).contiguous(), bias_name)
+
+        self.make_node(op_type, inputs=inputs, outputs=outputs, name=name, domain=("com.microsoft" if skip else None), **layernorm_kwargs)
+        
+    def make_qk_norm_layernorms(self, reshape_output, norm_weight, weight_name, layernorm_name, output_name, layernorm_kwargs):
+        if self.layernorm_attrs["use_fp32_layernorms"]:
+            cast_in = f"{layernorm_name}/Cast"
+            self.make_node("Cast", inputs=[reshape_output], outputs=[cast_in], name=f"{layernorm_name}/CastIn", to=TensorProto.FLOAT)
+            self.make_layernorm_node_and_tensors(
+                norm_weight,
+                weight_name,
+                None, None,
+                [cast_in, weight_name],
+                [f"{layernorm_name}/output/Cast"],
+                layernorm_name,
+                TensorProto.FLOAT,
+                False, True,
+                layernorm_kwargs
+            )
+            self.make_node("Cast", inputs=[f"{layernorm_name}/output/Cast"], outputs=[output_name], name=f"{layernorm_name}/CastOut", to=self.io_dtype)
+        else:
+            self.make_layernorm_node_and_tensors(
+                norm_weight,
+                weight_name,
+                None, None,
+                [reshape_output, weight_name],
+                [output_name],
+                layernorm_name,
+                self.io_dtype,
+                False, True,
+                layernorm_kwargs
+            )
+
     def make_qk_norm(self, layer_id, attention):
         # Make subgraph to compute SimplifiedLayerNorm after Q and K MatMuls in attention:
         #
@@ -1369,19 +1397,7 @@ def make_qk_norm(self, layer_id, attention):
         q_weight_name = f"model.layers.{layer_id}.attn.q_norm.layernorm.weight"
         q_layernorm_output = f"{q_layernorm_name}/output_0"
 
-        if "gemma-3" in self.model_name_or_path and self.io_dtype == TensorProto.BFLOAT16:
-            q_layernorm_cast_in = f"{q_layernorm_name}/Cast"
-            self.make_node("Cast", inputs=[q_reshape_1_output], outputs=[q_layernorm_cast_in], name=f"{q_layernorm_name}/CastIn", to=TensorProto.FLOAT)
-
-            q_layernorm_weight_fp32 = (attention.q_norm.weight.detach().cpu().to(self.to_torch_dtype[TensorProto.FLOAT]) + self.layernorm_attrs["add_offset"]).contiguous()
-            self.make_external_tensor(q_layernorm_weight_fp32, q_weight_name)
-
-            self.make_node("SimplifiedLayerNormalization", inputs=[q_layernorm_cast_in, q_weight_name], outputs=[f"{q_layernorm_name}/output/Cast"], name=q_layernorm_name, **layernorm_kwargs)
-            self.make_node("Cast", inputs=[f"{q_layernorm_name}/output/Cast"], outputs=[q_layernorm_output], name=f"{q_layernorm_name}/CastOut", to=self.io_dtype)
-        else:
-            self.make_external_tensor((attention.q_norm.weight.detach().cpu().to(self.to_torch_dtype[self.io_dtype]) + self.layernorm_attrs["add_offset"]).contiguous(), q_weight_name)
-            self.make_node("SimplifiedLayerNormalization", inputs=[q_reshape_1_output, q_weight_name], outputs=[q_layernorm_output], name=q_layernorm_name, **layernorm_kwargs)
-       
+        self.make_qk_norm_layernorms(reshape_output=q_reshape_1_output, norm_weight=attention.q_norm.weight, weight_name=q_weight_name, layernorm_name=q_layernorm_name, output_name=q_layernorm_output, layernorm_kwargs=layernorm_kwargs)
         self.make_value_info(q_layernorm_output, dtype=self.io_dtype, shape=['batch_size', 'sequence_length * num_attention_heads', self.head_size])
 
         # Reshape Q path after LayerNorm from Bx(SxN)xH to BxSxD
@@ -1399,19 +1415,7 @@ def make_qk_norm(self, layer_id, attention):
         k_weight_name = f"model.layers.{layer_id}.attn.k_norm.layernorm.weight"
         k_layernorm_output = f"{k_layernorm_name}/output_0"
 
-        if "gemma-3" in self.model_name_or_path and self.io_dtype == TensorProto.BFLOAT16:
-            k_layernorm_cast_in = f"{k_layernorm_name}/Cast"
-            self.make_node("Cast", inputs=[k_reshape_1_output], outputs=[k_layernorm_cast_in], name=f"{k_layernorm_name}/CastIn", to=TensorProto.FLOAT)
-
-            k_layernorm_weight_fp32 = (attention.k_norm.weight.detach().cpu().to(self.to_torch_dtype[TensorProto.FLOAT]) + self.layernorm_attrs["add_offset"]).contiguous()
-            self.make_external_tensor(k_layernorm_weight_fp32, k_weight_name)
-
-            self.make_node("SimplifiedLayerNormalization", inputs=[k_layernorm_cast_in, k_weight_name], outputs=[f"{k_layernorm_name}/output/Cast"], name=k_layernorm_name, **layernorm_kwargs)
-            self.make_node("Cast", inputs=[f"{k_layernorm_name}/output/Cast"], outputs=[k_layernorm_output], name=f"{k_layernorm_name}/CastOut", to=self.io_dtype)
-        else:
-            self.make_external_tensor((attention.k_norm.weight.detach().cpu().to(self.to_torch_dtype[self.io_dtype]) + self.layernorm_attrs["add_offset"]).contiguous(), k_weight_name)
-            self.make_node("SimplifiedLayerNormalization", inputs=[k_reshape_1_output, k_weight_name], outputs=[k_layernorm_output], name=k_layernorm_name, **layernorm_kwargs)                
-            
+        self.make_qk_norm_layernorms(reshape_output=k_reshape_1_output, norm_weight=attention.k_norm.weight, weight_name=k_weight_name, layernorm_name=k_layernorm_name, output_name=k_layernorm_output, layernorm_kwargs=layernorm_kwargs)
         self.make_value_info(k_layernorm_output, dtype=self.io_dtype, shape=['batch_size', 'sequence_length * num_key_value_heads', self.head_size])
 
         # Reshape K path after LayerNorm from Bx(SxN)xH to BxSxD

From 30077d08155d057bddf6efeb574abccc6ca14d0c Mon Sep 17 00:00:00 2001
From: Nenad Banfic <nebanfic@microsoft.com>
Date: Thu, 8 May 2025 00:53:50 +0000
Subject: [PATCH 19/26] Make sure bias is not used before being defined

---
 src/python/py/models/builder.py | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index e90599f0d2..c5a4d764ed 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -1055,10 +1055,11 @@ def make_layernorm_default(self, layer_id, layernorm, skip, simple, location):
         root_input = self.layernorm_attrs["root_input"]
         skip_input = self.layernorm_attrs["skip_input"]
 
+        bias_name = f"model.layers.{layer_id}.{location}_layernorm.bias"
         weight = f"model.layers.{layer_id}.{location}_layernorm.weight"
         inputs = [root_input, skip_input, weight] if skip else [root_input, weight]
         if not simple:
-            inputs.append(bias)
+            inputs.append(bias_name)
 
         name = f"/model/layers.{layer_id}/{location}_layernorm/{'Skip' if skip else ''}LayerNorm"
         kwargs = {"epsilon": self.layernorm_attrs["epsilon"]}
@@ -1071,7 +1072,6 @@ def make_layernorm_default(self, layer_id, layernorm, skip, simple, location):
             output_0 = "hidden_states"
         outputs = [output_0, "", "", output_3] if skip and not self.layernorm_attrs["last_layernorm"] else [output_0]
         
-        bias_name = f"model.layers.{layer_id}.{location}_layernorm.bias"
         bias = getattr(layernorm, "bias", None)
 
         self.make_layernorm_node_and_tensors(layernorm.weight, weight, bias, bias_name, inputs, outputs, name, TensorProto.FLOAT, skip, simple, kwargs)
@@ -1105,10 +1105,10 @@ def make_layernorm_with_casting_to_fp32(self, layer_id, layernorm, skip, simple,
             self.make_value_info(casted_skip, TensorProto.FLOAT, shape=["batch_size", "sequence_length", self.hidden_size],)
 
         weight = f"model.layers.{layer_id}.{location}_layernorm.weight"
-       
+        bias_name = f"model.layers.{layer_id}.{location}_layernorm.bias"
         inputs = [root_input, skip_input, weight] if skip else [root_input, weight]
         if not simple:
-            inputs.append(bias)
+            inputs.append(bias_name)
 
         name = f"/model/layers.{layer_id}/{location}_layernorm/{'Skip' if skip else ''}LayerNorm"
         kwargs = {"epsilon": self.layernorm_attrs["epsilon"]}
@@ -1123,7 +1123,6 @@ def make_layernorm_with_casting_to_fp32(self, layer_id, layernorm, skip, simple,
             raw_out0 = "hidden_states"
         outputs = [raw_out0, "", "", raw_out3] if skip and not self.layernorm_attrs["last_layernorm"] else [raw_out0]
 
-        bias_name = f"model.layers.{layer_id}.{location}_layernorm.bias"
         bias = getattr(layernorm, "bias", None)
         self.make_layernorm_node_and_tensors(layernorm.weight, weight, bias, bias_name, inputs, outputs, name, TensorProto.FLOAT, skip, simple, kwargs)
 

From 577e4b3b79021983db20ea617ee6ca0e82c7ffa6 Mon Sep 17 00:00:00 2001
From: Kunal Vaishnavi <kvaishnavi@microsoft.com>
Date: Thu, 8 May 2025 08:31:22 +0000
Subject: [PATCH 20/26] Redesign and refactor how casting is done

---
 src/python/py/models/builder.py | 332 ++++++++++++++++----------------
 1 file changed, 168 insertions(+), 164 deletions(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index c5a4d764ed..f6f74da27b 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -13,7 +13,7 @@
 from packaging import version
 
 if version.parse(ort_version) > version.parse("1.21.1"):
-    from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer, QuantFormat
+    from onnxruntime.quantization.matmul_nbits_quantizer import KQuantWeightOnlyQuantConfig, MatMulNBitsQuantizer, QuantFormat, RTNWeightOnlyQuantConfig
 else:
     from onnxruntime.quantization.matmul_4bits_quantizer import MatMul4BitsQuantizer as MatMulNBitsQuantizer, QuantFormat
 
@@ -100,14 +100,19 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
 
         # Map output names to their types and shapes
         self.output_names = ["logits"]
-        self.outputs_attrs = {"logits_type": self.io_dtype}
+        self.output_types = {
+            "hidden_states": self.io_dtype,                                                                      # For standard models where you want to remove the language modeling head from the model (note that `hidden_states` is written this way to match Hugging Face format)
+            "logits": self.io_dtype,                                                                             # For standard models
+            "present.key": self.io_dtype,                                                                        # For standard models (note that `present.key` is written this way to match Hugging Face format)
+            "present.value": self.io_dtype,                                                                      # For standard models (note that `present.value` is written this way to match Hugging Face format)
+        }
+        self.output_shapes = {
+            "hidden_states": ["batch_size", "sequence_length", self.hidden_size],                                # For standard models where you want to remove the language modeling head from the model (note that `hidden_states` is written this way to match Hugging Face format)
+            "logits": ["batch_size", "sequence_length", self.vocab_size],                                        # For standard models
+            "present.key": ["batch_size", self.num_kv_heads, "total_sequence_length", self.head_size],           # For standard models (note that `present.key` is written this way to match Hugging Face format)
+            "present.value": ["batch_size", self.num_kv_heads, "total_sequence_length", self.head_size],         # For standard models (note that `present.value` is written this way to match Hugging Face format)
+        }
         self.make_outputs_init()
-        self.exclude_lm_head = extra_options.get("exclude_lm_head", False)
-        self.include_hidden_states = extra_options.get("include_hidden_states", False)
-        if self.exclude_lm_head:
-            self.output_names = [name.replace("logits", "hidden_states") for name in self.output_names]
-        elif self.include_hidden_states:
-            self.output_names = ["hidden_states"] + self.output_names
 
         # Store names of nodes already created
         self.node_names = set()
@@ -172,6 +177,7 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
             "output_3": "",             # Output 3 for SkipLayerNorm
             "add_offset": 0,            # Offset value for LayerNorm weight
             "epsilon": epsilon,         # Epsilon value to avoid `sqrt(0)` in LayerNorm
+            "use_fp32": False,          # Use float32 precision to compute LayerNorm
         }
 
         # MatMul-specific variables
@@ -309,18 +315,13 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
 
 
     def make_outputs_init(self):
-        self.output_types = {
-            "hidden_states": self.io_dtype,                                                                      # For standard models where you want to remove the language modeling head from the model (note that `hidden_states` is written this way to match Hugging Face format)
-            "logits": self.outputs_attrs["logits_type"],                                                                               # For standard models
-            "present.key": self.io_dtype,                                                                        # For standard models (note that `present.key` is written this way to match Hugging Face format)
-            "present.value": self.io_dtype,                                                                      # For standard models (note that `present.value` is written this way to match Hugging Face format)
-        }
-        self.output_shapes = {
-            "hidden_states": ["batch_size", "sequence_length", self.hidden_size],                                # For standard models where you want to remove the language modeling head from the model (note that `hidden_states` is written this way to match Hugging Face format)
-            "logits": ["batch_size", "sequence_length", self.vocab_size],                                        # For standard models
-            "present.key": ["batch_size", self.num_kv_heads, "total_sequence_length", self.head_size],           # For standard models (note that `present.key` is written this way to match Hugging Face format)
-            "present.value": ["batch_size", self.num_kv_heads, "total_sequence_length", self.head_size],         # For standard models (note that `present.value` is written this way to match Hugging Face format)
-        }
+        self.exclude_lm_head = self.extra_options.get("exclude_lm_head", False)
+        self.include_hidden_states = self.extra_options.get("include_hidden_states", False)
+
+        if self.exclude_lm_head:
+            self.output_names = [name.replace("logits", "hidden_states") for name in self.output_names]
+        elif self.include_hidden_states:
+            self.output_names = ["hidden_states"] + self.output_names
 
     def make_attention_init(self):
         valid_gqa_configurations = [
@@ -494,10 +495,8 @@ def save_model(self, out_dir):
     def make_int4_algo_config(self, quant_method):
         int4_algo_config = None
         if quant_method == "rtn":
-            from onnxruntime.quantization.matmul_nbits_quantizer import RTNWeightOnlyQuantConfig
             int4_algo_config = RTNWeightOnlyQuantConfig()
         elif quant_method in ["k_quant_mixed", "k_quant_last"]:
-            from onnxruntime.quantization.matmul_nbits_quantizer import KQuantWeightOnlyQuantConfig
             if quant_method == "k_quant_mixed":
                 # k_quant_mixed is from llama.cpp.
                 # Reference: https://github.com/ggml-org/llama.cpp/blob/36667c8edcded08063ed51c7d57e9e086bbfc903/src/llama-quant.cpp#L136
@@ -821,15 +820,11 @@ def make_matmul_float(self, matmul, name, root_input, **kwargs):
 
         last_dim = matmul.weight.shape[0]
         output = "logits" if kwargs.get("logits", False) else f"{name}/output_0"
-
-        self.make_matmul_float_node_and_value_info(name, [root_input, weight], output, last_dim)
+        self.make_node("MatMul", inputs=[root_input, weight], outputs=[output], name=name)
+        self.make_value_info(output, self.io_dtype, shape=['batch_size', 'sequence_length', last_dim])
 
         return name
 
-    def make_matmul_float_node_and_value_info(self, name, inputs, output, last_dim):
-        self.make_node("MatMul", inputs=inputs, outputs=[output], name=name)
-        self.make_value_info(output, self.io_dtype, shape=['batch_size', 'sequence_length', last_dim])
-
     def make_matmul_int4(self, matmul, basename, root_input, **kwargs):
         if not hasattr(matmul, "qweight"):
             # TODO: quantize weights, then save new MatMul weights for onnx model
@@ -1051,103 +1046,115 @@ def make_embedding(self, embedding):
         self.layernorm_attrs["root_input"] = layernorm_attrs_value
         self.layernorm_attrs["skip_input"] = layernorm_attrs_value
 
-    def make_layernorm_default(self, layer_id, layernorm, skip, simple, location):
+    def make_layernorm(self, layer_id, layernorm, skip, simple, location):
         root_input = self.layernorm_attrs["root_input"]
         skip_input = self.layernorm_attrs["skip_input"]
 
-        bias_name = f"model.layers.{layer_id}.{location}_layernorm.bias"
+        # Get precision types to use
+        old_torch_dtype = self.to_torch_dtype[self.io_dtype]
+        old_io_dtype = self.io_dtype
+        new_torch_dtype = torch.float32 if self.layernorm_attrs["use_fp32"] else self.to_torch_dtype[self.io_dtype]
+        new_io_dtype = self.to_onnx_dtype[new_torch_dtype]
+        cast = (old_torch_dtype != new_torch_dtype)
+
+        # Create weight and bias tensors
         weight = f"model.layers.{layer_id}.{location}_layernorm.weight"
+        self.make_external_tensor(layernorm.weight.detach().cpu().to(new_torch_dtype).contiguous() + self.layernorm_attrs["add_offset"], weight)
+        bias = f"model.layers.{layer_id}.{location}_layernorm.bias"
+        if not simple:
+            self.make_external_tensor(layernorm.bias.detach().cpu().to(new_torch_dtype).contiguous(), bias)
+
+        # Create input names for op
         inputs = [root_input, skip_input, weight] if skip else [root_input, weight]
         if not simple:
-            inputs.append(bias_name)
+            inputs.append(bias)
 
         name = f"/model/layers.{layer_id}/{location}_layernorm/{'Skip' if skip else ''}LayerNorm"
+        op_type = f"{'Skip' if skip else ''}{'Simplified' if simple else ''}LayerNormalization"
         kwargs = {"epsilon": self.layernorm_attrs["epsilon"]}
         if not skip:
             kwargs.update({"axis": -1, "stash_type": 1})
 
+        # Create output names for op
         output_0 = f"/model/layers.{layer_id}/{location}_layernorm/output_0"
         output_3 = f"/model/layers.{layer_id}/{location}_layernorm/output_3"
         if self.layernorm_attrs["last_layernorm"] and (self.include_hidden_states or self.exclude_lm_head):
             output_0 = "hidden_states"
         outputs = [output_0, "", "", output_3] if skip and not self.layernorm_attrs["last_layernorm"] else [output_0]
-        
-        bias = getattr(layernorm, "bias", None)
 
-        self.make_layernorm_node_and_tensors(layernorm.weight, weight, bias, bias_name, inputs, outputs, name, TensorProto.FLOAT, skip, simple, kwargs)
-        self.make_value_info(output_0, self.io_dtype, shape=['batch_size', 'sequence_length', self.hidden_size])
+        # Create Cast nodes for inputs and outputs if old_dtype != new_dtype
+        if cast:
+            inputs, outputs = self.make_layernorm_casts(name, inputs, outputs, old_io_dtype, new_io_dtype)
+
+        # Make op and its shape
+        self.make_node(op_type, inputs=inputs, outputs=outputs, name=name, domain=("com.microsoft" if skip else None), **kwargs)
+        self.make_value_info(outputs[0], new_io_dtype, shape=['batch_size', 'sequence_length', self.hidden_size])
         if skip and not self.layernorm_attrs["last_layernorm"]:
-            self.make_value_info(output_3, self.io_dtype, shape=['batch_size', 'sequence_length', self.hidden_size])
+            self.make_value_info(outputs[3], new_io_dtype, shape=['batch_size', 'sequence_length', self.hidden_size])
 
+        # Update LayerNorm attributes
         self.layernorm_attrs["output_0"] = output_0
         if skip and not self.layernorm_attrs["last_layernorm"]:
             self.layernorm_attrs["output_3"] = output_3
 
+            # Assign output 3 of current SkipLayerNorm as root input to next SkipLayerNorm
             self.layernorm_attrs["root_input"] = output_3
 
-    def make_layernorm_with_casting_to_fp32(self, layer_id, layernorm, skip, simple, location):
-        root_input = self.layernorm_attrs["root_input"]
-        skip_input = self.layernorm_attrs["skip_input"]
-
-        skip_path = "skip" if skip else "noskip"
-
-        casted_root = f"{root_input}.{layer_id}.{skip_path}root/Cast"
-        self.make_node("Cast", inputs=[root_input], outputs=[casted_root], name=f"/model/layers.{layer_id}/{location}_layernorm/CastInput0", to=TensorProto.FLOAT)
-        root_input = casted_root
-
-        self.make_value_info(casted_root, TensorProto.FLOAT, shape=["batch_size", "sequence_length", self.hidden_size],)
+    def make_layernorm_casts(self, name, inputs, outputs, old_dtype, new_dtype):
+        # Name = name of original LayerNorm op as if the cast nodes did not exist
+        # Inputs = inputs into the original LayerNorm op as if the cast nodes did not exist
+        # Outputs = outputs from the original LayerNorm op as if the cast nodes did not exist
+        def get_shape_of_value_info(target_name):
+            for value_info in self.value_infos:
+                if value_info.name == target_name:
+                    shape = []
+                    for dim in value_info.type.tensor_type.shape.dim:
+                        if dim.HasField("dim_value"):
+                            shape.append(dim.dim_value)
+                        elif dim.HasField("dim_param"):
+                            shape.append(dim.dim_param)
+                        else:
+                            shape.append(None)
+                    return shape
+
+        # Save original inputs and outputs
+        skip = len(inputs) > 2  # [root_input, skip_input, weight] vs. [root_input, weight]
+        root_input = inputs[0]
+        skip_input = inputs[1] if skip else None
+        output_0 = outputs[0]
+        output_3 = outputs[3] if skip and not self.layernorm_attrs["last_layernorm"] else None
+
+        # Cast root_input
+        root_input_cast_name = f"{name}/root_input/Cast"
+        root_input_cast_output = f"{root_input_cast_name}/output_0"
+        self.make_node("Cast", inputs=[root_input], outputs=[root_input_cast_output], name=root_input_cast_name, to=new_dtype)
+        self.make_value_info(root_input_cast_output, new_dtype, shape=get_shape_of_value_info(root_input))
+        inputs[0] = root_input_cast_output
 
         if skip:
-            casted_skip = f"{skip_input}.{layer_id}.{skip_path}skip/Cast"
-            self.make_node("Cast", inputs=[skip_input], outputs=[casted_skip], name=f"/model/layers.{layer_id}/{location}_layernorm/CastInput3", to=TensorProto.FLOAT)
-            skip_input = casted_skip
-
-            self.make_value_info(casted_skip, TensorProto.FLOAT, shape=["batch_size", "sequence_length", self.hidden_size],)
+            # Cast skip_input
+            skip_input_cast_name = f"{name}/skip_input/Cast"
+            skip_input_cast_output = f"{skip_input_cast_name}/output_0"
+            self.make_node("Cast", inputs=[skip_input], outputs=[skip_input_cast_output], name=skip_input_cast_name, to=new_dtype)
+            self.make_value_info(skip_input_cast_output, new_dtype, shape=get_shape_of_value_info(skip_input))
+            inputs[1] = skip_input_cast_output
+
+        # Cast output_0
+        output_0_cast_name = f"{name}/output_0/Cast"
+        output_0_cast_output = f"{output_0_cast_name}/output_0"
+        self.make_node("Cast", inputs=[output_0_cast_output], outputs=[output_0], name=output_0_cast_name, to=old_dtype)
+        self.make_value_info(output_0, old_dtype, shape=get_shape_of_value_info(root_input))
+        outputs[0] = output_0_cast_output
 
-        weight = f"model.layers.{layer_id}.{location}_layernorm.weight"
-        bias_name = f"model.layers.{layer_id}.{location}_layernorm.bias"
-        inputs = [root_input, skip_input, weight] if skip else [root_input, weight]
-        if not simple:
-            inputs.append(bias_name)
-
-        name = f"/model/layers.{layer_id}/{location}_layernorm/{'Skip' if skip else ''}LayerNorm"
-        kwargs = {"epsilon": self.layernorm_attrs["epsilon"]}
-        if not skip:
-            kwargs.update({"axis": -1, "stash_type": 1})
-
-        output_0 = f"/model/layers.{layer_id}/{location}_layernorm/output_0"
-        output_3 = f"/model/layers.{layer_id}/{location}_layernorm/output_3"
-        raw_out0 = f"{output_0}/Cast"
-        raw_out3 = f"{output_3}/Cast"
-        if self.layernorm_attrs["last_layernorm"] and (self.include_hidden_states or self.exclude_lm_head):
-            raw_out0 = "hidden_states"
-        outputs = [raw_out0, "", "", raw_out3] if skip and not self.layernorm_attrs["last_layernorm"] else [raw_out0]
-
-        bias = getattr(layernorm, "bias", None)
-        self.make_layernorm_node_and_tensors(layernorm.weight, weight, bias, bias_name, inputs, outputs, name, TensorProto.FLOAT, skip, simple, kwargs)
-
-        self.make_value_info(raw_out0, TensorProto.FLOAT, shape=['batch_size', 'sequence_length', self.hidden_size])
         if skip and not self.layernorm_attrs["last_layernorm"]:
-            self.make_value_info(raw_out3, TensorProto.FLOAT, shape=['batch_size', 'sequence_length', self.hidden_size])
-
-        if not(skip and not self.layernorm_attrs["last_layernorm"]):
-            raw_out3 = None
-
-        self.make_node("Cast", inputs=[raw_out0], outputs=[output_0], name=f"{name}/CastOutput0", to=self.io_dtype,)
-        self.make_value_info(output_0, self.io_dtype, shape=["batch_size", "sequence_length", self.hidden_size],)
-        self.layernorm_attrs["output_0"] = output_0
-
-        if raw_out3 is not None:
-            self.make_node("Cast", inputs=[raw_out3], outputs=[output_3], name=f"{name}/CastOutput3", to=self.io_dtype,)
-            self.make_value_info(output_3, self.io_dtype, shape=["batch_size", "sequence_length", self.hidden_size],)
-            self.layernorm_attrs["output_3"] = output_3
-            self.layernorm_attrs["root_input"] = output_3
+            # Cast output_3
+            output_3_cast_name = f"{name}/output_3/Cast"
+            output_3_cast_output = f"{output_3_cast_name}/output_3"
+            self.make_node("Cast", inputs=[output_3_cast_output], outputs=[output_3], name=output_3_cast_name, to=old_dtype)
+            self.make_value_info(output_3, old_dtype, shape=get_shape_of_value_info(root_input))
+            outputs[3] = output_3_cast_output
 
-    def make_layernorm(self, layer_id, layernorm, skip, simple, location):
-        if self.layernorm_attrs["use_fp32_layernorms"]:
-            self.make_layernorm_with_casting_to_fp32(layer_id, layernorm, skip, simple, location)
-        else:
-            self.make_layernorm_default(layer_id, layernorm, skip, simple, location)
+        return inputs, outputs
 
     def make_mscale_su(self, mscale):
         if mscale <= 1.0:
@@ -1334,44 +1341,6 @@ def make_rotary_embedding_multi_cache(self, **kwargs):
         self.make_value_info(cos_cache_name, self.io_dtype, shape=["max_sequence_length", "head_dim / 2"])
         self.make_value_info(sin_cache_name, self.io_dtype, shape=["max_sequence_length", "head_dim / 2"])
 
-    def make_layernorm_node_and_tensors(self, weight, weight_name, bias, bias_name, inputs, outputs, name, tensor_type, skip, simple, layernorm_kwargs):
-        op_type = f"{'Skip' if skip else ''}{'Simplified' if simple else ''}LayerNormalization"
-        self.make_external_tensor((weight.detach().cpu().to(self.to_torch_dtype[tensor_type]) + self.layernorm_attrs["add_offset"]).contiguous(), weight_name)
-        
-        if not simple:
-            self.make_external_tensor(bias.detach().to(self.to_torch_dtype[TensorProto.FLOAT]).contiguous(), bias_name)
-
-        self.make_node(op_type, inputs=inputs, outputs=outputs, name=name, domain=("com.microsoft" if skip else None), **layernorm_kwargs)
-        
-    def make_qk_norm_layernorms(self, reshape_output, norm_weight, weight_name, layernorm_name, output_name, layernorm_kwargs):
-        if self.layernorm_attrs["use_fp32_layernorms"]:
-            cast_in = f"{layernorm_name}/Cast"
-            self.make_node("Cast", inputs=[reshape_output], outputs=[cast_in], name=f"{layernorm_name}/CastIn", to=TensorProto.FLOAT)
-            self.make_layernorm_node_and_tensors(
-                norm_weight,
-                weight_name,
-                None, None,
-                [cast_in, weight_name],
-                [f"{layernorm_name}/output/Cast"],
-                layernorm_name,
-                TensorProto.FLOAT,
-                False, True,
-                layernorm_kwargs
-            )
-            self.make_node("Cast", inputs=[f"{layernorm_name}/output/Cast"], outputs=[output_name], name=f"{layernorm_name}/CastOut", to=self.io_dtype)
-        else:
-            self.make_layernorm_node_and_tensors(
-                norm_weight,
-                weight_name,
-                None, None,
-                [reshape_output, weight_name],
-                [output_name],
-                layernorm_name,
-                self.io_dtype,
-                False, True,
-                layernorm_kwargs
-            )
-
     def make_qk_norm(self, layer_id, attention):
         # Make subgraph to compute SimplifiedLayerNorm after Q and K MatMuls in attention:
         #
@@ -1383,8 +1352,13 @@ def make_qk_norm(self, layer_id, attention):
         #          |
         #       Reshape (BxSxD)
 
-        # Save kwargs shared by LayerNorm ops
+        # Save kwargs shared by LayerNorm ops and precision types to use
         layernorm_kwargs = {"epsilon": self.layernorm_attrs["epsilon"], "axis": -1, "stash_type": 1}
+        old_torch_dtype = self.to_torch_dtype[self.io_dtype]
+        old_io_dtype = self.io_dtype
+        new_torch_dtype = torch.float32 if self.layernorm_attrs["use_fp32"] else self.to_torch_dtype[self.io_dtype]
+        new_io_dtype = self.to_onnx_dtype[new_torch_dtype]
+        cast = (old_torch_dtype != new_torch_dtype)
 
         # Reshape Q MatMul from BxSxD to Bx(SxN)xH before LayerNorm
         q_reshape_1_name = f"/model/layers.{layer_id}/attn/q_norm/Reshape_1"
@@ -1392,12 +1366,20 @@ def make_qk_norm(self, layer_id, attention):
         q_reshape_1_output = f"{q_reshape_1_name}/output_0"
         self.make_reshape(q_reshape_1_name, q_reshape_1_inputs, dtype=self.io_dtype, shape=['batch_size', 'sequence_length * num_attention_heads', self.head_size])
 
+        # Make Q LayerNorm
         q_layernorm_name = f"/model/layers.{layer_id}/attn/q_norm/SimplifiedLayerNormalization"
         q_weight_name = f"model.layers.{layer_id}.attn.q_norm.layernorm.weight"
         q_layernorm_output = f"{q_layernorm_name}/output_0"
+        self.make_external_tensor((attention.q_norm.weight.detach().cpu().to(new_torch_dtype) + self.layernorm_attrs["add_offset"]).contiguous(), q_weight_name)
+
+        # Create Cast nodes for inputs and outputs if old_dtype != new_dtype
+        q_layernorm_inputs = [q_reshape_1_output, q_weight_name]
+        q_layernorm_outputs = [q_layernorm_output]
+        if cast:
+            q_layernorm_inputs, q_layernorm_outputs = self.make_layernorm_casts(q_layernorm_name, q_layernorm_inputs, q_layernorm_outputs, old_io_dtype, new_io_dtype)
 
-        self.make_qk_norm_layernorms(reshape_output=q_reshape_1_output, norm_weight=attention.q_norm.weight, weight_name=q_weight_name, layernorm_name=q_layernorm_name, output_name=q_layernorm_output, layernorm_kwargs=layernorm_kwargs)
-        self.make_value_info(q_layernorm_output, dtype=self.io_dtype, shape=['batch_size', 'sequence_length * num_attention_heads', self.head_size])
+        self.make_node("SimplifiedLayerNormalization", inputs=q_layernorm_inputs, outputs=q_layernorm_outputs, name=q_layernorm_name, **layernorm_kwargs)
+        self.make_value_info(q_layernorm_outputs[0], dtype=new_io_dtype, shape=['batch_size', 'sequence_length * num_attention_heads', self.head_size])
 
         # Reshape Q path after LayerNorm from Bx(SxN)xH to BxSxD
         q_reshape_2_name = f"/model/layers.{layer_id}/attn/q_norm/Reshape_2"
@@ -1410,12 +1392,20 @@ def make_qk_norm(self, layer_id, attention):
         k_reshape_1_output = f"{k_reshape_1_name}/output_0"
         self.make_reshape(k_reshape_1_name, k_reshape_1_inputs, dtype=self.io_dtype, shape=['batch_size', 'sequence_length * num_key_value_heads', self.head_size])
 
+        # Make K LayerNorm
         k_layernorm_name = f"/model/layers.{layer_id}/attn/k_norm/SimplifiedLayerNormalization"
         k_weight_name = f"model.layers.{layer_id}.attn.k_norm.layernorm.weight"
         k_layernorm_output = f"{k_layernorm_name}/output_0"
+        self.make_external_tensor((attention.k_norm.weight.detach().cpu().to(new_torch_dtype) + self.layernorm_attrs["add_offset"]).contiguous(), k_weight_name)
+
+        # Create Cast nodes for inputs and outputs if old_dtype != new_dtype
+        k_layernorm_inputs = [k_reshape_1_output, k_weight_name]
+        k_layernorm_outputs = [k_layernorm_output]
+        if cast:
+            k_layernorm_inputs, k_layernorm_outputs = self.make_layernorm_casts(k_layernorm_name, k_layernorm_inputs, k_layernorm_outputs, old_io_dtype, new_io_dtype)
 
-        self.make_qk_norm_layernorms(reshape_output=k_reshape_1_output, norm_weight=attention.k_norm.weight, weight_name=k_weight_name, layernorm_name=k_layernorm_name, output_name=k_layernorm_output, layernorm_kwargs=layernorm_kwargs)
-        self.make_value_info(k_layernorm_output, dtype=self.io_dtype, shape=['batch_size', 'sequence_length * num_key_value_heads', self.head_size])
+        self.make_node("SimplifiedLayerNormalization", inputs=k_layernorm_inputs, outputs=k_layernorm_outputs, name=k_layernorm_name, **layernorm_kwargs)
+        self.make_value_info(k_layernorm_outputs[0], dtype=new_io_dtype, shape=['batch_size', 'sequence_length * num_key_value_heads', self.head_size])
 
         # Reshape K path after LayerNorm from Bx(SxN)xH to BxSxD
         k_reshape_2_name = f"/model/layers.{layer_id}/attn/k_norm/Reshape_2"
@@ -1711,10 +1701,10 @@ def make_attention(self, layer_id, attention, root_input, **kwargs):
             q_matmul_name = self.make_matmul(attention.q_proj, q_matmul_basename, root_input)
             self.attention_attrs["q_path"] = f"{q_matmul_name}/output_0"
             k_matmul_basename = f"/model/layers.{layer_id}/attn/k_proj/MatMul"
-            k_matmul_name = self.make_matmul(attention.k_proj, k_matmul_basename, root_input, first_cast=False)
+            k_matmul_name = self.make_matmul(attention.k_proj, k_matmul_basename, root_input)
             self.attention_attrs["k_path"] = f"{k_matmul_name}/output_0"
             v_matmul_basename = f"/model/layers.{layer_id}/attn/v_proj/MatMul"
-            v_matmul_name = self.make_matmul(attention.v_proj, v_matmul_basename, root_input, first_cast=False)
+            v_matmul_name = self.make_matmul(attention.v_proj, v_matmul_basename, root_input)
             self.attention_attrs["v_path"] = f"{v_matmul_name}/output_0"
 
         # Make Add nodes (if bias exists)
@@ -1980,7 +1970,7 @@ def make_mlp_proj(self, layer_id, mlp, root_input):
 
         # Make Up proj nodes
         up_matmul_basename = f"/model/layers.{layer_id}/mlp/up_proj/MatMul"
-        up_matmul_name = self.make_matmul(mlp.up_proj, up_matmul_basename, root_input, first_cast=False)
+        up_matmul_name = self.make_matmul(mlp.up_proj, up_matmul_basename, root_input)
         up_name = up_matmul_name
         if up_bias_exists:
             up_add_name = f"/model/layers.{layer_id}/mlp/up_proj/Add"
@@ -2238,21 +2228,26 @@ def make_lm_head(self, lm_head):
         scale_exists = self.lm_head_attrs["scale"] != 1
         mask_exists = self.lm_head_attrs["mask"] is not None
         softcap_exists = self.lm_head_attrs["softcap"] != 0.0
+        cast_exists = (self.io_dtype != self.output_types["logits"])
+
+        # List order matters here. It should match the order of the below if condition checks.
+        # Add new checks to the end of the list and after the below if condition checks.
+        exists_checks = [bias_exists, scale_exists, mask_exists, softcap_exists, cast_exists]
 
         matmul_basename = "/lm_head/MatMul"
         root_input = self.layernorm_attrs["output_0"]
-        matmul_name = self.make_matmul(lm_head, matmul_basename, root_input, logits=not(bias_exists or scale_exists or mask_exists or softcap_exists))
+        matmul_name = self.make_matmul(lm_head, matmul_basename, root_input, logits=not any(exists_checks))
         lm_name = matmul_name
 
         if bias_exists:
             add_name = "/lm_head/Add"
-            self.make_add_bias(lm_head.bias.detach().cpu(), add_name, root_input=f"{lm_name}/output_0", logits=not(scale_exists or mask_exists or softcap_exists))
+            self.make_add_bias(lm_head.bias.detach().cpu(), add_name, root_input=f"{lm_name}/output_0", logits=not any(exists_checks[1:]))
             lm_name = add_name
 
         if scale_exists:
             mul_name = "/lm_head/Mul"
             mul_inputs = [f"{lm_name}/output_0", f"/model/constants/{self.to_str_dtype[self.io_dtype]}/0D/{self.lm_head_attrs['scale']}"]
-            mul_output = "logits" if not(mask_exists or softcap_exists) else f"{mul_name}/output_0"
+            mul_output = "logits" if not any(exists_checks[2:]) else f"{mul_name}/output_0"
             self.make_node('Mul', inputs=mul_inputs, outputs=[mul_output], name=mul_name)
             self.make_value_info(mul_output, self.io_dtype, shape=['batch_size', 'sequence_length', self.vocab_size])
             lm_name = mul_name
@@ -2264,7 +2259,7 @@ def make_lm_head(self, lm_head):
 
             where_name = "/lm_head/Where"
             where_inputs = [logits_mask_name, f"/model/constants/{self.to_str_dtype[self.io_dtype]}/0D/{torch.finfo(self.to_torch_dtype[self.io_dtype]).min}", f"{lm_name}/output_0"]
-            where_output = "logits" if not softcap_exists else f"{where_name}/output_0"
+            where_output = "logits" if not any(exists_checks[3:]) else f"{where_name}/output_0"
             self.make_node('Where', inputs=where_inputs, outputs=[where_output], name=where_name)
             self.make_value_info(where_output, self.io_dtype, shape=['batch_size', 'sequence_length', self.vocab_size])
             lm_name = where_name
@@ -2280,10 +2275,17 @@ def make_lm_head(self, lm_head):
 
             mul_name = "/lm_head/softcap/Mul"
             mul_inputs = [f"{tanh_name}/output_0", f"/model/constants/{self.to_str_dtype[self.io_dtype]}/0D/{self.lm_head_attrs['softcap']}"]
-            mul_output = "logits"
+            mul_output = "logits" if not any(exists_checks[4:]) else f"{mul_name}/output_0"
             self.make_node('Mul', inputs=mul_inputs, outputs=[mul_output], name=mul_name)
             self.make_value_info(mul_output, self.io_dtype, shape=['batch_size', 'sequence_length', self.vocab_size])
-            lm_head = mul_name
+            lm_name = mul_name
+
+        if cast_exists:
+            # Add final cast from io_dtype to logits_dtype
+            cast_name = "/lm_head/Cast"
+            cast_output = "logits"
+            self.make_node('Cast', inputs=[f"{lm_name}/output_0"], outputs=[cast_output], to=self.output_types['logits'])
+            self.make_value_info(cast_output, self.output_types['logits'], shape=['batch_size', 'sequence_length', self.vocab_size])
 
     def make_layer(self, layer_id, layer):
         # Each LLM decoder layer is typically defined as:
@@ -2314,7 +2316,7 @@ def make_model(self, input_path):
                 from onnxruntime_genai.models.gguf_model import GGUFModel
             model = GGUFModel.from_pretrained(self.model_type, input_path, self.head_size, self.hidden_size, self.intermediate_size, self.num_attn_heads, self.num_kv_heads, self.vocab_size)
             self.layernorm_attrs["add_offset"] = 0  # add offset already done for GGUF models
-            self.layernorm_attrs["use_fp32_layernorms"] = 0
+
         elif self.quant_type is not None:
             # Load quantized PyTorch model
             try:
@@ -2324,6 +2326,7 @@ def make_model(self, input_path):
             q_size = self.num_attn_heads * self.head_size
             kv_size = self.num_kv_heads * self.head_size
             model = QuantModel.from_pretrained(self.quant_type, input_path, self.quant_attrs, q_size, kv_size, self.intermediate_size, self.num_layers)
+
         else:
             # Load PyTorch model
             extra_kwargs = {"num_hidden_layers": self.num_layers} if "num_hidden_layers" in self.extra_options else {}
@@ -2902,9 +2905,15 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
 class Gemma2Model(GemmaModel):
     def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
         super().__init__(config, io_dtype, onnx_dtype, ep, cache_dir, extra_options)
+        self.layernorm_attrs["use_fp32"] = True
         self.attention_attrs["scale"] = config.query_pre_attn_scalar ** -0.5
         self.is_local = lambda layer_id: layer_id % 2 == 1
 
+    def make_outputs_init(self):
+        # Always use float32 logits to improve accuracy
+        self.output_types["logits"] = TensorProto.FLOAT
+        super().make_outputs_init()
+
     def make_layer(self, layer_id, layer):
         # Gemma2 decoder layer is typically defined as:
         # input_layernorm --> attention --> post_attention_layernorm --> pre_ffn_layernorm --> MLP --> post_ffn_layernorm
@@ -3343,27 +3352,12 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
         self.is_local = lambda layer_id: bool((layer_id + 1) % config.sliding_window_pattern)
         self.rope_local_theta = config.rope_local_base_freq
         self.make_rotary_embedding_multi_cache()
-        self.layernorm_attrs["use_fp32_layernorms"] = 1 if self.io_dtype == TensorProto.BFLOAT16 else 0        
 
     def make_attention_init(self):
         self.attention_attrs["q_norm"] = True
         self.attention_attrs["k_norm"] = True
         super().make_attention_init()
 
-    def make_outputs_init(self):
-        self.outputs_attrs["logits_type"] = TensorProto.FLOAT
-        super().make_outputs_init()
-
-    def make_matmul_float_node_and_value_info(self, name, inputs, output, last_dim):
-        if output == "logits":
-            cast_name = name + "/Cast"
-            cast_name_input = name + "/Cast"
-            self.make_node("MatMul", inputs=inputs, outputs=[cast_name_input], name=name)
-            self.make_node("Cast", inputs=[cast_name_input], outputs=[output], name=cast_name, to=TensorProto.FLOAT)
-            self.make_value_info(output, TensorProto.FLOAT, shape=['batch_size', 'sequence_length', last_dim])
-        else:
-            super().make_matmul_float_node_and_value_info(name, inputs, output, last_dim)
-
     def make_rotary_embedding_multi_cache(self):
         self.cos_cache_global_name, self.sin_cache_global_name = "cos_cache_global", "sin_cache_global"
         super().make_rotary_embedding_caches(cos_cache_name=self.cos_cache_global_name, sin_cache_name=self.sin_cache_global_name)
@@ -3487,8 +3481,14 @@ def create_model(model_name, input_path, output_dir, precision, execution_provid
         elif config.architectures[0] == "GemmaForCausalLM":
             onnx_model = GemmaModel(config, io_dtype, precision, execution_provider, cache_dir, extra_options)
         elif config.architectures[0] == "Gemma2ForCausalLM":
+            if precision == "fp16":
+                print("WARNING: This model loses accuracy with float16 precision. Setting `--precision bf16` by default.")
+                precision = "bf16"
             onnx_model = Gemma2Model(config, io_dtype, precision, execution_provider, cache_dir, extra_options)
         elif config.architectures[0] == "Gemma3ForCausalLM":
+            if precision == "fp16":
+                print("WARNING: This model loses accuracy with float16 precision. Setting `--precision bf16` by default.")
+                precision = "bf16"
             onnx_model = Gemma3Model(config, io_dtype, precision, execution_provider, cache_dir, extra_options)
             onnx_model.model_type = "gemma3_text"
         elif config.architectures[0] == "Gemma3ForConditionalGeneration":
@@ -3497,7 +3497,11 @@ def create_model(model_name, input_path, output_dir, precision, execution_provid
             for key in text_config:
                 if not hasattr(config, key):
                     setattr(config, key, getattr(text_config, key))
+
             # extra_options["exclude_embeds"] = True
+            if precision == "fp16":
+                print("WARNING: This model loses accuracy with float16 precision. Setting `--precision bf16` by default.")
+                precision = "bf16"
             onnx_model = Gemma3Model(config, io_dtype, precision, execution_provider, cache_dir, extra_options)
             onnx_model.model_type = "gemma3_text"
         elif config.architectures[0] == "GraniteForCausalLM":
@@ -3518,7 +3522,7 @@ def create_model(model_name, input_path, output_dir, precision, execution_provid
             onnx_model = Phi3MiniLongRoPEModel(config, io_dtype, precision, execution_provider, cache_dir, extra_options)
         elif config.architectures[0] == "PhiMoEForCausalLM" and config.max_position_embeddings != config.original_max_position_embeddings:
             print("WARNING: This model only works for CUDA currently because `MoE` is only supported for CUDA in ONNX Runtime. Setting `--execution_provider cuda` by default.")
-            print("WARNING: This model currently only supports quantized version. Setting `--precision int4` by default.")
+            print("WARNING: This model currently only supports the quantized version. Setting `--precision int4` by default.")
             execution_provider = "cuda"
             precision = "int4"
             onnx_model = Phi3MoELongRoPEModel(config, io_dtype, precision, execution_provider, cache_dir, extra_options)
@@ -3678,4 +3682,4 @@ def get_args():
 if __name__ == '__main__':
     args = get_args()
     extra_options = parse_extra_options(args.extra_options)
-    create_model(args.model_name, args.input, args.output, args.precision, args.execution_provider, args.cache_dir, **extra_options)
\ No newline at end of file
+    create_model(args.model_name, args.input, args.output, args.precision, args.execution_provider, args.cache_dir, **extra_options)

From a6bb4654b161a21d217630e411a11727ed3e542d Mon Sep 17 00:00:00 2001
From: Kunal Vaishnavi <kvaishnavi@microsoft.com>
Date: Thu, 8 May 2025 08:47:53 +0000
Subject: [PATCH 21/26] Fix embedding check to be more generic

---
 src/python/py/models/builder.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index f6f74da27b..a7b1a3af29 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -1246,7 +1246,6 @@ def make_rotary_embedding(self, name, root_input, **kwargs):
         num_heads = self.num_kv_heads if "k_rotary" in name else self.num_attn_heads
 
         inputs = [root_input, kwargs.pop("position_ids"), cos_cache_name, sin_cache_name]
-
         output = f"{name}/output_0"
         self.make_node(
             "RotaryEmbedding", inputs=inputs, outputs=[output], name=name, domain="com.microsoft",
@@ -2339,9 +2338,9 @@ def make_model(self, input_path):
         # Loop through model and map each module to ONNX/ORT ops
         self.layer_id = 0
         for module in model.modules():
-            if isinstance(module, torch.nn.Embedding) or (hasattr(model, "embedding") and module == model.embedding):
+            if (isinstance(module, torch.nn.Embedding) and module.weight.shape[0] == self.vocab_size) or (hasattr(model, "embedding") and module == model.embedding):
                 # Checks (Hugging Face logic) or (GGUF logic)
-                if not self.exclude_embeds and module == model.language_model.model.embed_tokens:
+                if not self.exclude_embeds:
                     # Embedding layer
                     print("Reading embedding layer")
                     self.make_embedding(module.weight.detach().cpu())

From 01b2a4977757b7808e5694e138e686408d476424 Mon Sep 17 00:00:00 2001
From: Kunal Vaishnavi <kvaishnavi@microsoft.com>
Date: Thu, 8 May 2025 08:49:05 +0000
Subject: [PATCH 22/26] Add missing cast name

---
 src/python/py/models/builder.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index a7b1a3af29..d5c2a8016a 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -2283,7 +2283,7 @@ def make_lm_head(self, lm_head):
             # Add final cast from io_dtype to logits_dtype
             cast_name = "/lm_head/Cast"
             cast_output = "logits"
-            self.make_node('Cast', inputs=[f"{lm_name}/output_0"], outputs=[cast_output], to=self.output_types['logits'])
+            self.make_node('Cast', inputs=[f"{lm_name}/output_0"], outputs=[cast_output], name=cast_name, to=self.output_types['logits'])
             self.make_value_info(cast_output, self.output_types['logits'], shape=['batch_size', 'sequence_length', self.vocab_size])
 
     def make_layer(self, layer_id, layer):

From b6d8b94ba3be90b843b2d05cd452d977e00afd96 Mon Sep 17 00:00:00 2001
From: Kunal Vaishnavi <kvaishnavi@microsoft.com>
Date: Thu, 8 May 2025 08:56:43 +0000
Subject: [PATCH 23/26] Revert quantization import consolidation until CIs are
 updated

---
 src/python/py/models/builder.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index d5c2a8016a..6f975d3dc2 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -13,7 +13,7 @@
 from packaging import version
 
 if version.parse(ort_version) > version.parse("1.21.1"):
-    from onnxruntime.quantization.matmul_nbits_quantizer import KQuantWeightOnlyQuantConfig, MatMulNBitsQuantizer, QuantFormat, RTNWeightOnlyQuantConfig
+    from onnxruntime.quantization.matmul_nbits_quantizer import MatMulNBitsQuantizer, QuantFormat
 else:
     from onnxruntime.quantization.matmul_4bits_quantizer import MatMul4BitsQuantizer as MatMulNBitsQuantizer, QuantFormat
 
@@ -495,8 +495,10 @@ def save_model(self, out_dir):
     def make_int4_algo_config(self, quant_method):
         int4_algo_config = None
         if quant_method == "rtn":
+            from onnxruntime.quantization.matmul_nbits_quantizer import RTNWeightOnlyQuantConfig
             int4_algo_config = RTNWeightOnlyQuantConfig()
         elif quant_method in ["k_quant_mixed", "k_quant_last"]:
+            from onnxruntime.quantization.matmul_nbits_quantizer import KQuantWeightOnlyQuantConfig
             if quant_method == "k_quant_mixed":
                 # k_quant_mixed is from llama.cpp.
                 # Reference: https://github.com/ggml-org/llama.cpp/blob/36667c8edcded08063ed51c7d57e9e086bbfc903/src/llama-quant.cpp#L136

From 3eb173dbf7898939f5ab1ee3faf78ead697a1a36 Mon Sep 17 00:00:00 2001
From: Kunal Vaishnavi <kvaishnavi@microsoft.com>
Date: Thu, 8 May 2025 09:00:19 +0000
Subject: [PATCH 24/26] Remove unneeded parentheses

---
 src/python/py/models/builder.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index 6f975d3dc2..63f4d9f304 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -1057,7 +1057,7 @@ def make_layernorm(self, layer_id, layernorm, skip, simple, location):
         old_io_dtype = self.io_dtype
         new_torch_dtype = torch.float32 if self.layernorm_attrs["use_fp32"] else self.to_torch_dtype[self.io_dtype]
         new_io_dtype = self.to_onnx_dtype[new_torch_dtype]
-        cast = (old_torch_dtype != new_torch_dtype)
+        cast = old_torch_dtype != new_torch_dtype
 
         # Create weight and bias tensors
         weight = f"model.layers.{layer_id}.{location}_layernorm.weight"
@@ -1359,7 +1359,7 @@ def make_qk_norm(self, layer_id, attention):
         old_io_dtype = self.io_dtype
         new_torch_dtype = torch.float32 if self.layernorm_attrs["use_fp32"] else self.to_torch_dtype[self.io_dtype]
         new_io_dtype = self.to_onnx_dtype[new_torch_dtype]
-        cast = (old_torch_dtype != new_torch_dtype)
+        cast = old_torch_dtype != new_torch_dtype
 
         # Reshape Q MatMul from BxSxD to Bx(SxN)xH before LayerNorm
         q_reshape_1_name = f"/model/layers.{layer_id}/attn/q_norm/Reshape_1"
@@ -2229,7 +2229,7 @@ def make_lm_head(self, lm_head):
         scale_exists = self.lm_head_attrs["scale"] != 1
         mask_exists = self.lm_head_attrs["mask"] is not None
         softcap_exists = self.lm_head_attrs["softcap"] != 0.0
-        cast_exists = (self.io_dtype != self.output_types["logits"])
+        cast_exists = self.io_dtype != self.output_types["logits"]
 
         # List order matters here. It should match the order of the below if condition checks.
         # Add new checks to the end of the list and after the below if condition checks.

From f8f6e027e8f77eed53ca70fcb87fdc51d17da0ac Mon Sep 17 00:00:00 2001
From: Kunal Vaishnavi <kvaishnavi@microsoft.com>
Date: Thu, 8 May 2025 23:43:46 +0000
Subject: [PATCH 25/26] Add more granular control over LayerNorm casts

---
 src/python/py/models/builder.py | 85 ++++++++++++++++++++++++---------
 1 file changed, 63 insertions(+), 22 deletions(-)

diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index 63f4d9f304..e1e86fa4e4 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -177,7 +177,13 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
             "output_3": "",             # Output 3 for SkipLayerNorm
             "add_offset": 0,            # Offset value for LayerNorm weight
             "epsilon": epsilon,         # Epsilon value to avoid `sqrt(0)` in LayerNorm
-            "use_fp32": False,          # Use float32 precision to compute LayerNorm
+            "cast": {                   # Casting LayerNorm-specific variables
+                "use_fp32": False,      # Use float32 precision to compute LayerNorm
+                "root_input": False,    # Cast root_input
+                "skip_input": False,    # Cast skip_input
+                "output_0": False,      # Cast output_0
+                "output_3": False,      # Cast output_3
+            }
         }
 
         # MatMul-specific variables
@@ -1055,7 +1061,7 @@ def make_layernorm(self, layer_id, layernorm, skip, simple, location):
         # Get precision types to use
         old_torch_dtype = self.to_torch_dtype[self.io_dtype]
         old_io_dtype = self.io_dtype
-        new_torch_dtype = torch.float32 if self.layernorm_attrs["use_fp32"] else self.to_torch_dtype[self.io_dtype]
+        new_torch_dtype = torch.float32 if self.layernorm_attrs["cast"]["use_fp32"] else self.to_torch_dtype[self.io_dtype]
         new_io_dtype = self.to_onnx_dtype[new_torch_dtype]
         cast = old_torch_dtype != new_torch_dtype
 
@@ -1126,14 +1132,15 @@ def get_shape_of_value_info(target_name):
         output_0 = outputs[0]
         output_3 = outputs[3] if skip and not self.layernorm_attrs["last_layernorm"] else None
 
-        # Cast root_input
-        root_input_cast_name = f"{name}/root_input/Cast"
-        root_input_cast_output = f"{root_input_cast_name}/output_0"
-        self.make_node("Cast", inputs=[root_input], outputs=[root_input_cast_output], name=root_input_cast_name, to=new_dtype)
-        self.make_value_info(root_input_cast_output, new_dtype, shape=get_shape_of_value_info(root_input))
-        inputs[0] = root_input_cast_output
+        if self.layernorm_attrs["cast"]["root_input"]:
+            # Cast root_input
+            root_input_cast_name = f"{name}/root_input/Cast"
+            root_input_cast_output = f"{root_input_cast_name}/output_0"
+            self.make_node("Cast", inputs=[root_input], outputs=[root_input_cast_output], name=root_input_cast_name, to=new_dtype)
+            self.make_value_info(root_input_cast_output, new_dtype, shape=get_shape_of_value_info(root_input))
+            inputs[0] = root_input_cast_output
 
-        if skip:
+        if skip and self.layernorm_attrs["cast"]["skip_input"]:
             # Cast skip_input
             skip_input_cast_name = f"{name}/skip_input/Cast"
             skip_input_cast_output = f"{skip_input_cast_name}/output_0"
@@ -1141,14 +1148,15 @@ def get_shape_of_value_info(target_name):
             self.make_value_info(skip_input_cast_output, new_dtype, shape=get_shape_of_value_info(skip_input))
             inputs[1] = skip_input_cast_output
 
-        # Cast output_0
-        output_0_cast_name = f"{name}/output_0/Cast"
-        output_0_cast_output = f"{output_0_cast_name}/output_0"
-        self.make_node("Cast", inputs=[output_0_cast_output], outputs=[output_0], name=output_0_cast_name, to=old_dtype)
-        self.make_value_info(output_0, old_dtype, shape=get_shape_of_value_info(root_input))
-        outputs[0] = output_0_cast_output
+        if self.layernorm_attrs["cast"]["output_0"]:
+            # Cast output_0
+            output_0_cast_name = f"{name}/output_0/Cast"
+            output_0_cast_output = f"{output_0_cast_name}/output_0"
+            self.make_node("Cast", inputs=[output_0_cast_output], outputs=[output_0], name=output_0_cast_name, to=old_dtype)
+            self.make_value_info(output_0, old_dtype, shape=get_shape_of_value_info(root_input))
+            outputs[0] = output_0_cast_output
 
-        if skip and not self.layernorm_attrs["last_layernorm"]:
+        if skip and not self.layernorm_attrs["last_layernorm"] and self.layernorm_attrs["cast"]["output_3"]:
             # Cast output_3
             output_3_cast_name = f"{name}/output_3/Cast"
             output_3_cast_output = f"{output_3_cast_name}/output_3"
@@ -1357,7 +1365,7 @@ def make_qk_norm(self, layer_id, attention):
         layernorm_kwargs = {"epsilon": self.layernorm_attrs["epsilon"], "axis": -1, "stash_type": 1}
         old_torch_dtype = self.to_torch_dtype[self.io_dtype]
         old_io_dtype = self.io_dtype
-        new_torch_dtype = torch.float32 if self.layernorm_attrs["use_fp32"] else self.to_torch_dtype[self.io_dtype]
+        new_torch_dtype = torch.float32 if self.layernorm_attrs["cast"]["use_fp32"] else self.to_torch_dtype[self.io_dtype]
         new_io_dtype = self.to_onnx_dtype[new_torch_dtype]
         cast = old_torch_dtype != new_torch_dtype
 
@@ -2906,7 +2914,11 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
 class Gemma2Model(GemmaModel):
     def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
         super().__init__(config, io_dtype, onnx_dtype, ep, cache_dir, extra_options)
-        self.layernorm_attrs["use_fp32"] = True
+        self.layernorm_attrs["cast"]["use_fp32"] = True
+        self.layernorm_attrs["cast"]["root_input"] = True
+        self.layernorm_attrs["cast"]["skip_input"] = False
+        self.layernorm_attrs["cast"]["output_0"] = True
+        self.layernorm_attrs["cast"]["output_3"] = False
         self.attention_attrs["scale"] = config.query_pre_attn_scalar ** -0.5
         self.is_local = lambda layer_id: layer_id % 2 == 1
 
@@ -2915,30 +2927,59 @@ def make_outputs_init(self):
         self.output_types["logits"] = TensorProto.FLOAT
         super().make_outputs_init()
 
+    def make_layernorm(self, layer_id, layernorm, skip, simple, location):
+        if "final_norm" in location:
+            # Set cast for final LayerNorm since it is a special case and not covered in `make_layer`
+            self.layernorm_attrs["cast"]["root_input"] = False
+        super().make_layernorm(layer_id, layernorm, skip, simple, location)
+
     def make_layer(self, layer_id, layer):
         # Gemma2 decoder layer is typically defined as:
         # input_layernorm --> attention --> post_attention_layernorm --> pre_ffn_layernorm --> MLP --> post_ffn_layernorm
+
+        # Adjust LayerNorm attributes because of extra LayerNorms inserted
+        # 1. Only cast root_input if the first layer of LayerNorms are being created
+        original_cast_root_input = self.layernorm_attrs["cast"]["root_input"]
+        self.layernorm_attrs["cast"]["root_input"] = self.layernorm_attrs["first_layernorm"]
         self.make_layernorm(layer_id, layer.input_layernorm, skip=not self.layernorm_attrs["first_layernorm"], simple=self.layernorm_attrs["simple"], location="input")
+        self.layernorm_attrs["cast"]["root_input"] = original_cast_root_input
+
         self.make_attention(layer_id, layer.self_attn, root_input=self.layernorm_attrs["output_0"])
 
-        # Temporarily set root_input for LayerNorm to skip_input for post_attention_layernorm
-        # Set skip_input to output of post_attention_layernorm
+        # Adjust LayerNorm attributes for extra LayerNorm to insert
+        # 1. Temporarily set root_input for LayerNorm to skip_input for post_attention_layernorm
+        # 2. Set skip_input to output of post_attention_layernorm
+        # 3. Do not cast outputs from post_attention_layernorm
         original_root_input = self.layernorm_attrs["root_input"]
+        original_cast_output_0 = self.layernorm_attrs["cast"]["output_0"]
         self.layernorm_attrs["root_input"] = self.layernorm_attrs["skip_input"]
+        self.layernorm_attrs["cast"]["output_0"] = False
         self.make_layernorm(layer_id, layer.post_attention_layernorm, skip=False, simple=self.layernorm_attrs["simple"], location="post_attention")
         self.layernorm_attrs["root_input"] = original_root_input
         self.layernorm_attrs["skip_input"] = self.layernorm_attrs["output_0"]
+        self.layernorm_attrs["cast"]["output_0"] = original_cast_output_0
 
+        # Adjust LayerNorm attributes because of extra LayerNorms inserted
+        # 1. Only cast root_input if the first layer of LayerNorms are being created
+        original_cast_root_input = self.layernorm_attrs["cast"]["root_input"]
+        self.layernorm_attrs["cast"]["root_input"] = self.layernorm_attrs["first_layernorm"]
         self.make_layernorm(layer_id, layer.pre_feedforward_layernorm, skip=True, simple=self.layernorm_attrs["simple"], location="pre_feedforward")
+        self.layernorm_attrs["cast"]["root_input"] = original_cast_root_input
+
         self.make_mlp(layer_id, layer.mlp, root_input=self.layernorm_attrs["output_0"])
 
-        # Temporarily set root_input for LayerNorm to skip_input for post_feedforward_layernorm
-        # Set skip_input to output of post_ffn_layernorm
+        # Adjust LayerNorm attributes for extra LayerNorm to insert
+        # 1. Temporarily set root_input for LayerNorm to skip_input for post_feedforward_layernorm
+        # 2. Set skip_input to output of post_feedforward_layernorm
+        # 3. Do not cast outputs from post_feedforward_layernorm
         original_root_input = self.layernorm_attrs["root_input"]
+        original_cast_output_0 = self.layernorm_attrs["cast"]["output_0"]
         self.layernorm_attrs["root_input"] = self.layernorm_attrs["skip_input"]
+        self.layernorm_attrs["cast"]["output_0"] = False
         self.make_layernorm(layer_id, layer.post_feedforward_layernorm, skip=False, simple=self.layernorm_attrs["simple"], location="post_feedforward")
         self.layernorm_attrs["root_input"] = original_root_input
         self.layernorm_attrs["skip_input"] = self.layernorm_attrs["output_0"]
+        self.layernorm_attrs["cast"]["output_0"] = original_cast_output_0
 
         self.layernorm_attrs["first_layernorm"] = False
         if layer_id == self.num_layers - 1:

From cdd273c2c8ec628c93885f600e1ddc77c76edd7a Mon Sep 17 00:00:00 2001
From: Kunal Vaishnavi <kvaishnavi@microsoft.com>
Date: Fri, 9 May 2025 00:27:00 +0000
Subject: [PATCH 26/26] Use same chat template for Gemma models

---
 examples/python/model-chat.py   | 2 +-
 examples/python/model-qa.py     | 4 ++--
 src/python/py/models/builder.py | 3 ---
 3 files changed, 3 insertions(+), 6 deletions(-)

diff --git a/examples/python/model-chat.py b/examples/python/model-chat.py
index ce2d4c6ccd..ef64e9d15b 100644
--- a/examples/python/model-chat.py
+++ b/examples/python/model-chat.py
@@ -56,7 +56,7 @@ def main(args):
             print("Using Chat Template for LLAMA 3, if you are using LLAMA  2 please pass the argument --chat_template '{input} [/INST]')")
         elif model_type.startswith("qwen2"):
             args.chat_template = '<|im_start|>user\n{input}<|im_end|>\n<|im_start|>assistant\n'
-        elif model_type == "gemma3_text":
+        elif model_type.startswith("gemma"):
             args.chat_template = '<start_of_turn>user\n{system_prompt}{input}<end_of_turn>\n<start_of_turn>model\n'
         else:
             raise ValueError(f"Chat Template for model type {model_type} is not known. Please provide chat template using --chat_template")
diff --git a/examples/python/model-qa.py b/examples/python/model-qa.py
index 0e1d6675db..c9f2fa5d06 100644
--- a/examples/python/model-qa.py
+++ b/examples/python/model-qa.py
@@ -56,7 +56,7 @@ def main(args):
             print("Using Chat Template for LLAMA 3, if you are using LLAMA  2 please pass the argument --chat_template '{input} [/INST]')")
         elif model_type.startswith("qwen2"):
             args.chat_template = '{system_prompt}<|im_start|>user\n{input}<|im_end|>\n<|im_start|>assistant\n'
-        elif model_type == "gemma3_text":
+        elif model_type.startswith("gemma"):
             args.chat_template = '<start_of_turn>user\n{system_prompt}{input}<end_of_turn>\n<start_of_turn>model\n'
         else:
             raise ValueError(f"Chat Template for model type {model_type} is not known. Please provide chat template using --chat_template")
@@ -75,7 +75,7 @@ def main(args):
             print("Using System Prompt for LLAMA 3, if you are using LLAMA  2 please pass the argument --system_prompt '<s>[INST] <<SYS>>\\n{args.system_prompt}\\n<</SYS>>')")
         elif model_type.startswith("qwen2"):
             system_prompt = f"<|im_start|>system\n{args.system_prompt}<|im_end|>\n"
-        elif model_type == "gemma3_text":
+        elif model_type.startswith("gemma"):
             system_prompt = f"{args.system_prompt}"
         else:
             system_prompt = args.system_prompt
diff --git a/src/python/py/models/builder.py b/src/python/py/models/builder.py
index e1e86fa4e4..3569898cca 100644
--- a/src/python/py/models/builder.py
+++ b/src/python/py/models/builder.py
@@ -3539,13 +3539,10 @@ def create_model(model_name, input_path, output_dir, precision, execution_provid
             for key in text_config:
                 if not hasattr(config, key):
                     setattr(config, key, getattr(text_config, key))
-
-            # extra_options["exclude_embeds"] = True
             if precision == "fp16":
                 print("WARNING: This model loses accuracy with float16 precision. Setting `--precision bf16` by default.")
                 precision = "bf16"
             onnx_model = Gemma3Model(config, io_dtype, precision, execution_provider, cache_dir, extra_options)
-            onnx_model.model_type = "gemma3_text"
         elif config.architectures[0] == "GraniteForCausalLM":
             onnx_model = GraniteModel(config, io_dtype, precision, execution_provider, cache_dir, extra_options)
         elif config.architectures[0] == "LlamaForCausalLM":