NvTensorRtRtx EP option in GenAI - model builder

src/python/py/models/builder.py

@@ -74,6 +74,7 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
             },
             "dml": {},
             "webgpu": {},
+            "NvTensorRtRtx": {},
         }
 
         # Map input names to their types and shapes
@@ -310,7 +311,7 @@ def __init__(self, config, io_dtype, onnx_dtype, ep, cache_dir, extra_options):
                 "nodes_to_exclude": extra_options.get("int4_nodes_to_exclude", []),
                 "algo_config": int4_algo_config,
             },
-            "use_qdq": extra_options.get("use_qdq", False),
+            "use_qdq": True if self.ep == "NvTensorRtRtx" else extra_options.get("use_qdq", False),
         }
         if self.quant_type is not None:
             # Create quantized attributes from quantization config
@@ -327,6 +328,7 @@ def make_attention_init(self):
             ("dml", TensorProto.FLOAT16),
             ("webgpu", TensorProto.FLOAT16),
             ("webgpu", TensorProto.FLOAT),
+            ("NvTensorRtRtx", TensorProto.FLOAT16),
         ]
         if (self.ep, self.io_dtype) in valid_gqa_configurations:
             # Change model settings for GroupQueryAttention
@@ -741,6 +743,23 @@ def make_reduce_max(self, name, inputs, dtype, shape):
         self.make_node("ReduceMax", inputs=inputs, outputs=[output], name=name, keepdims=False)
         self.make_value_info(output, dtype, shape=shape)
 
+    def make_reduce_mean(self, name, inputs, dtype, shape, axes=[-1], keepdims=False):
+        output = f"{name}/output_0"
+        if self.quant_attrs["use_qdq"]:
+            # Opset 18 uses axes as input[1]
+            inputs.append(f"/model/constants/TensorProto.INT64/1D/{','.join(map(str, axes))}")
+            self.make_node("ReduceMean", inputs=inputs, outputs=[output], name=name, keepdims=keepdims)
+            self.make_value_info(output, dtype, shape=shape)
+        else:
+            # Opset 17 uses axes as attribute
+            self.make_node("ReduceMean", inputs=inputs, outputs=[output], name=name, axes=axes, keepdims=keepdims)
+            self.make_value_info(output, dtype, shape=shape)
+
+    def make_sqrt(self, name, inputs, dtype, shape):
+        output = f"{name}/output_0"
+        self.make_node("Sqrt", inputs=inputs, outputs=[output], name=name)
+        self.make_value_info(output, dtype, shape=shape)
+
     def make_cast(self, name, root_input, dtype, shape):
         output = f"{name}/output_0"
         self.make_node("Cast", inputs=[root_input], outputs=[output], name=name, to=dtype)
@@ -1068,8 +1087,28 @@ def make_layernorm(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]
 
-        self.make_node(op_type, inputs=inputs, outputs=outputs, name=name, domain=("com.microsoft" if skip else None), **kwargs)
-        self.make_value_info(output_0, self.io_dtype, shape=['batch_size', 'sequence_length', self.hidden_size])
+        # NvTensorRtRtx EP doesn't support Skip/SimplifiedLayerNormalization, so we fallback to primitive ops
+        if self.ep == "NvTensorRtRtx":
+            layer_norn_basename = f"/model/layers.{layer_id}/{location}_layernorm"
+            if op_type == "SimplifiedLayerNormalization":
+                output_0 = f"{layer_norn_basename}/simplified_layer_norm/Mul_1/output_0"
+                outputs = [output_0]
+                self.make_simplified_layer_norm(layer_norn_basename, root_input, weight, shape=['batch_size', 'sequence_length', self.hidden_size])
+            elif op_type == "SkipSimplifiedLayerNormalization":
+                output_0 = f"{layer_norn_basename}/skip_simplified_layer_norm/simplified_layer_norm/Mul_1/output_0"
+                output_3 = f"{layer_norn_basename}/skip_simplified_layer_norm/Add/output_0"
+                self.make_skip_simplified_layer_norm(layer_norn_basename, root_input, skip_input, weight, shape=['batch_size', 'sequence_length', self.hidden_size])
+            elif op_type == "SkipLayerNormalization":
+                output_0 = f"{layer_norn_basename}/skip_layer_norm/LayerNormalization/output_0"
+                output_3 = f"{layer_norn_basename}/skip_layer_norm/Add/output_0"
+                self.make_skip_layer_norm(layer_norn_basename, root_input, skip_input, weight, bias, shape=['batch_size', 'sequence_length', self.hidden_size])
+            else: # LayerNormalization
+                self.make_node(op_type, inputs=inputs, outputs=outputs, name=name, **kwargs)
+                self.make_value_info(output_0, self.io_dtype, shape=['batch_size', 'sequence_length', self.hidden_size])
+        else:
+            self.make_node(op_type, inputs=inputs, outputs=outputs, name=name, domain=("com.microsoft" if skip else None), **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])
 
@@ -1265,6 +1304,105 @@ 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"])
 
+    # This expansion of contrib-op can be updated / depricated in future.
+    def make_skip_simplified_layer_norm(self, basename, root_input, skip_input, weight_name, shape, **kwargs):
+        #                          root_input         skip_input
+        #                              |                  |
+        #                              +------------------+
+        #                              |
+        #                             Add-------------> output (1)
+        #                              |
+        #                      SimplifiedLayerNorm----> output (0)
+        make_add_name = f"{basename}/skip_simplified_layer_norm/Add"
+        make_add_inputs = [root_input, skip_input]
+        self.make_add(make_add_name, make_add_inputs, dtype=self.io_dtype, shape=shape)
+
+        make_simplified_layer_norm_name = f"{basename}/skip_simplified_layer_norm"
+        self.make_simplified_layer_norm(make_simplified_layer_norm_name,  f"{make_add_name}/output_0", weight_name, shape=shape, **kwargs)
+
+    # This expansion contrib-op can be updated / depricated in future.
+    def make_skip_layer_norm(self, basename, root_input, skip_input, weight_name, bias_name, shape, **kwargs):
+        #                          root_input         skip_input
+        #                              |                  |
+        #                              +------------------+
+        #                              |
+        #                             Add-------------> output (1)
+        #                              |
+        #                      LayerNormalization-----> output (0)
+        make_add_name = f"{basename}/skip_layer_norm/Add"
+        make_add_inputs = [root_input, skip_input]
+        self.make_add(make_add_name, make_add_inputs, dtype=self.io_dtype, shape=shape)
+
+        kwargs = {"epsilon": self.layernorm_attrs["epsilon"]}
+        kwargs.update({"axis": -1, "stash_type": 1})
+
+        make_layer_norm_name = f"{basename}/skip_layer_norm/LayerNormalization"
+        inputs = [f"{make_add_name}/output_0", weight_name, bias_name]
+        outputs = [f"{make_layer_norm_name}/output_0"]
+        self.make_node("LayerNormalization", inputs=inputs, outputs=outputs, name=make_layer_norm_name, **kwargs)
+        self.make_value_info(f"{make_layer_norm_name}/output_0", self.io_dtype, shape=shape)
+
+    # This expansion contrib-op can be updated / depricated in future.
+    def make_simplified_layer_norm(self, basename, root_input, weight_name, shape, **kwargs):
+
+        #                            Cast (float32) - most calc happens in higher precision
+        #                              |
+        #                      +-------+-------+
+        #                      |               |
+        #                     Pow              |
+        #                      |               |
+        #                  ReduceMean          |
+        #                      |               |
+        #                     Add              |
+        #                      |               |
+        #                    Sqrt              |
+        #                      |               |
+        #                     Div              |
+        #                      |               |
+        #                      +-------+-------+
+        #                              |  
+        #                             Mul
+        #                              |
+        #                            Cast_1 (float16)
+        #                              |
+        #                            Mul_1
+
+        make_cast_name = f"{basename}/simplified_layer_norm/Cast"
+        self.make_cast(make_cast_name, root_input, TensorProto.FLOAT, shape=shape)
+
+        make_pow_name = f"{basename}/simplified_layer_norm/Pow"
+        make_pow_inputs = [f"{make_cast_name}/output_0", f"/model/constants/TensorProto.FLOAT/0D/2"]
+
+        self.make_node("Pow", inputs=make_pow_inputs, outputs=[f"{make_pow_name}/output_0"], name=make_pow_name, domain="")
+        self.make_value_info(f"{make_pow_name}/output_0", TensorProto.FLOAT, shape=shape)
+
+        make_reducemean_name = f"{basename}/simplified_layer_norm/ReduceMean"
+        make_reducemean_inputs = [f"{make_pow_name}/output_0"]
+        self.make_reduce_mean(make_reducemean_name, make_reducemean_inputs, TensorProto.FLOAT, keepdims=True, axes=[-1], shape=shape)
+
+        make_add_name = f"{basename}/simplified_layer_norm/Add"
+        make_add_inputs = [f"{make_reducemean_name}/output_0", f"/model/constants/TensorProto.FLOAT/0D/{self.layernorm_attrs['epsilon']}"]
+        self.make_add(make_add_name, make_add_inputs, TensorProto.FLOAT, shape=shape)
+
+        make_sqrt_name = f"{basename}/simplified_layer_norm/Sqrt"
+        make_sqrt_inputs = [f"{make_add_name}/output_0"]
+        self.make_sqrt(make_sqrt_name, make_sqrt_inputs, TensorProto.FLOAT, shape=shape)
+
+        make_div_name = f"{basename}/simplified_layer_norm/Div"
+        make_div_inputs = [f"/model/constants/TensorProto.FLOAT/0D/1", f"{make_sqrt_name}/output_0"]
+        self.make_div(make_div_name, make_div_inputs, TensorProto.FLOAT, shape=shape)
+
+        make_mul_name = f"{basename}/simplified_layer_norm/Mul"
+        make_mul_inputs = [f"{make_div_name}/output_0", f"{make_cast_name}/output_0"]
+        self.make_mul(make_mul_name, make_mul_inputs, TensorProto.FLOAT, shape=shape)
+
+        make_cast_1_name = f"{basename}/simplified_layer_norm/Cast_1"
+        self.make_cast(make_cast_1_name, f"{make_mul_name}/output_0", dtype=self.io_dtype, shape=shape)
+
+        make_mul_1_name = f"{basename}/simplified_layer_norm/Mul_1"
+        make_mul_1_inputs = [f"{make_cast_1_name}/output_0", weight_name]
+        self.make_mul(make_mul_1_name, make_mul_1_inputs, dtype=self.io_dtype, shape=shape)
+
     def make_qk_norm(self, layer_id, attention):
         # Make subgraph to compute SimplifiedLayerNorm after Q and K MatMuls in attention:
         #
@@ -2089,7 +2227,15 @@ def make_gelu(self, layer_id, root_input, activation):
         #        GeluAct
         gelu_name = f"/model/layers.{layer_id}/mlp/act_fn/{activation}"
         output = f"{gelu_name}/output_0"
-        self.make_node(activation, inputs=[root_input], outputs=[output], name=gelu_name, domain="com.microsoft")
+
+        # NvTensorRtRtx (Opset 21) uses standard "Gelu" replacing "Gelu" & "FastGelu" contrib ops, otherwise fallback to contrib ops
+        if activation == "Gelu" and self.ep == "NvTensorRtRtx":
+            self.make_node("Gelu", inputs=[root_input], outputs=[output], name=gelu_name, approximate="none") 
+        elif activation == "FastGelu" and self.ep == "NvTensorRtRtx":
+            self.make_node("Gelu", inputs=[root_input], outputs=[output], name=gelu_name, approximate="tanh")
+        else:
+            self.make_node(activation, inputs=[root_input], outputs=[output], name=gelu_name, domain="com.microsoft")
+
         self.make_value_info(output, self.io_dtype, shape=['batch_size', 'sequence_length', self.intermediate_size])
 
         return gelu_name
@@ -3469,7 +3615,7 @@ def get_args():
         "-e",
         "--execution_provider",
         required=True,
-        choices=["cpu", "cuda", "rocm", "dml", "webgpu"],
+        choices=["cpu", "cuda", "rocm", "dml", "webgpu", "NvTensorRtRtx"],
         help="Execution provider to target with precision of model (e.g. FP16 CUDA, INT4 CPU, INT4 WEBGPU)",
     )