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)
@@ -1052,23 +1071,30 @@ def make_layernorm(self, layer_id, layernorm, skip, simple, location):
         if not simple:
             self.make_external_tensor(layernorm.bias.detach().cpu().to(self.to_torch_dtype[self.io_dtype]).contiguous(), bias)
 
+        name = f"/model/layers.{layer_id}/{location}_layernorm/{'Skip' if skip else ''}LayerNorm"
+        # NvTensorRtRtx EP doesn't support Skip/SimplifiedLayerNormalization and SkipLayerNormalization, so we fallback to primitive ops
+        if self.ep == "NvTensorRtRtx" and (skip or simple):
+            self._make_layernormlization(name, skip, simple, root_input, skip_input, weight, bias)
+        else:
+            self.make_layernormlization(name, skip, simple, root_input, skip_input, weight, bias)
+
+    def make_layernormlization(self, basename, skip, simple, root_input, skip_input, weight, 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"
+        output_0 = f"{basename}/output_0"
+        output_3 = f"{basename}/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]
 
-        self.make_node(op_type, inputs=inputs, outputs=outputs, name=name, domain=("com.microsoft" if skip else None), **kwargs)
+        kwargs = {"epsilon": self.layernorm_attrs["epsilon"]}
+        if not skip:
+            kwargs.update({"axis": -1, "stash_type": 1})
+
+        op_type = f"{'Skip' if skip else ''}{'Simplified' if simple else ''}LayerNormalization"
+        self.make_node(op_type, inputs=inputs, outputs=outputs, name=basename, 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])
@@ -1081,6 +1107,17 @@ 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_layernormlization(self, basename, skip, simple, root_input, skip_input, weight, bias):
+        op_type = f"{'Skip' if skip else ''}{'Simplified' if simple else ''}LayerNormalization"
+        if op_type == "SimplifiedLayerNormalization":
+            self._make_simplified_layer_norm(basename, root_input, weight, shape=['batch_size', 'sequence_length', self.hidden_size])
+        elif op_type == "SkipSimplifiedLayerNormalization":
+            self._make_skip_simplified_layer_norm(basename, root_input, skip_input, weight, shape=['batch_size', 'sequence_length', self.hidden_size])
+        elif op_type == "SkipLayerNormalization":
+            self._make_skip_layer_norm(basename, root_input, skip_input, weight, bias, shape=['batch_size', 'sequence_length', self.hidden_size])
+        else:
+            raise ValueError(f"Invalid op_type: {op_type}")
+
     def make_mscale_su(self, mscale):
         if mscale <= 1.0:
             return 1.0
@@ -1265,6 +1302,132 @@ 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 / deprecated 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)
+
+        if not self.layernorm_attrs["last_layernorm"]:
+            self.make_value_info(f"{make_add_name}/output_0", self.io_dtype, shape=['batch_size', 'sequence_length', self.hidden_size])
+            self.layernorm_attrs["output_3"] = f"{make_add_name}/output_0"      
+            # Assign output 3 of current SkipLayerNorm as root input to next SkipLayerNorm
+            self.layernorm_attrs["root_input"] = f"{make_add_name}/output_0"
+
+    # 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]
+        output = f"{make_layer_norm_name}/output_0"
+
+        if self.layernorm_attrs["last_layernorm"] and (self.include_hidden_states or self.exclude_lm_head):
+            output = ["hidden_states"]
+
+        self.make_node("LayerNormalization", inputs=inputs, outputs=[output], name=make_layer_norm_name, **kwargs)
+        self.make_value_info(f"{make_layer_norm_name}/output_0", self.io_dtype, shape=shape)
+
+        self.layernorm_attrs["output_0"] = output
+        if not self.layernorm_attrs["last_layernorm"]:
+            self.make_value_info(f"{make_add_name}/output_0", self.io_dtype, shape=['batch_size', 'sequence_length', self.hidden_size])
+            self.layernorm_attrs["output_3"] = f"{make_add_name}/output_0"        
+            # Assign output 3 of current SkipLayerNorm as root input to next SkipLayerNorm
+            self.layernorm_attrs["root_input"] = f"{make_add_name}/output_0"
+
+    # 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]
+
+        make_mul_1_output = f"{make_mul_1_name}/output_0"
+        if self.layernorm_attrs["last_layernorm"] and (self.include_hidden_states or self.exclude_lm_head):
+            make_mul_1_output = "hidden_states"
+
+        self.make_node("Mul", inputs=make_mul_1_inputs, outputs=[make_mul_1_output], name=make_mul_1_name)
+        self.make_value_info(make_mul_1_output, dtype=self.io_dtype, shape=shape)
+
+        # Update LayerNorm attributes
+        self.layernorm_attrs["output_0"] = make_mul_1_output
+
+
     def make_qk_norm(self, layer_id, attention):
         # Make subgraph to compute SimplifiedLayerNorm after Q and K MatMuls in attention:
         #
@@ -2082,17 +2245,47 @@ def make_activation_with_mul(self, layer_id, root_input, activation, domain):
         return mul_act_name
 
     def make_gelu(self, layer_id, root_input, activation):
+        # NvTensorRtRtx (Opset 21) uses standard "Gelu" replacing "Gelu" & "FastGelu" contrib ops, otherwise fallback to contrib ops
+        if self.ep == "NvTensorRtRtx" and activation in ["Gelu", "FastGelu"]:
+            return self._make_gelu_op(layer_id, root_input, activation)
+        else:
+            return self.make_gelu_op(layer_id, root_input, activation)
+
+    def make_gelu_op(self, layer_id, root_input, activation):
         # Make nodes for this activation subgraph
         #
         #       root_input (Add)
         #           |
         #        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")
         self.make_value_info(output, self.io_dtype, shape=['batch_size', 'sequence_length', self.intermediate_size])
 
         return gelu_name
+
+    # This expansion of contrib-op can be updated / deprecated in future.
+    def _make_gelu_op(self, layer_id, root_input, activation):
+        # Make nodes for this activation subgraph
+        #
+        #       root_input (Add)
+        #           |
+        #        GeluAct
+        gelu_name = f"/model/layers.{layer_id}/mlp/act_fn/{activation}"
+        output = f"{gelu_name}/output_0"
+
+        # NvTensorRtRtx (Opset 21) uses standard "Gelu" replacing "Gelu" & "FastGelu" contrib ops, otherwise fallback to contrib ops
+        if activation == "Gelu":
+            self.make_node("Gelu", inputs=[root_input], outputs=[output], name=gelu_name, approximate="none") 
+        elif activation == "FastGelu":
+            self.make_node("Gelu", inputs=[root_input], outputs=[output], name=gelu_name, approximate="tanh")
+        else:
+            raise NotImplementedError(f"The {activation} activation function is not currently supported.")
+
+        self.make_value_info(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}"
@@ -3469,7 +3662,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)",
     )