test: Add L0_backend_onnxruntime test for enabling bfloat16 dtype in ONNXRuntime backend

docs/user_guide/model_configuration.md

@@ -326,22 +326,22 @@ If a model framework does not have an entry for a given datatype, then Triton do
 The sixth column, labeled "API", shows the corresponding datatype for the TRITONSERVER C API, TRITONBACKEND C API, HTTP/REST protocol and GRPC protocol.
 The last column shows the corresponding datatype for the Python numpy library.
 
-|Model Config  |TensorRT      |ONNX Runtime  |PyTorch  |API      |NumPy         |
-|--------------|--------------|--------------|---------|---------|--------------|
-|TYPE_BOOL     | kBOOL        |BOOL          |kBool    |BOOL     |bool          |
-|TYPE_UINT8    | kUINT8       |UINT8         |kByte    |UINT8    |uint8         |
-|TYPE_UINT16   |              |UINT16        |         |UINT16   |uint16        |
-|TYPE_UINT32   |              |UINT32        |         |UINT32   |uint32        |
-|TYPE_UINT64   |              |UINT64        |         |UINT64   |uint64        |
-|TYPE_INT8     | kINT8        |INT8          |kChar    |INT8     |int8          |
-|TYPE_INT16    |              |INT16         |kShort   |INT16    |int16         |
-|TYPE_INT32    | kINT32       |INT32         |kInt     |INT32    |int32         |
-|TYPE_INT64    | kINT64       |INT64         |kLong    |INT64    |int64         |
-|TYPE_FP16     | kHALF        |FLOAT16       |         |FP16     |float16       |
-|TYPE_FP32     | kFLOAT       |FLOAT         |kFloat   |FP32     |float32       |
-|TYPE_FP64     |              |DOUBLE        |kDouble  |FP64     |float64       |
-|TYPE_STRING   |              |STRING        |         |BYTES    |dtype(object) |
-|TYPE_BF16     | kBF16        |              |         |BF16     |              |
+|Model Config  |TensorRT      |ONNX Runtime  |PyTorch  |API      |NumPy              |
+|--------------|--------------|--------------|---------|---------|-------------------|
+|TYPE_BOOL     | kBOOL        |BOOL          |kBool    |BOOL     |bool               |
+|TYPE_UINT8    | kUINT8       |UINT8         |kByte    |UINT8    |uint8              |
+|TYPE_UINT16   |              |UINT16        |         |UINT16   |uint16             |
+|TYPE_UINT32   |              |UINT32        |         |UINT32   |uint32             |
+|TYPE_UINT64   |              |UINT64        |         |UINT64   |uint64             |
+|TYPE_INT8     | kINT8        |INT8          |kChar    |INT8     |int8               |
+|TYPE_INT16    |              |INT16         |kShort   |INT16    |int16              |
+|TYPE_INT32    | kINT32       |INT32         |kInt     |INT32    |int32              |
+|TYPE_INT64    | kINT64       |INT64         |kLong    |INT64    |int64              |
+|TYPE_FP16     | kHALF        |FLOAT16       |         |FP16     |float16            |
+|TYPE_FP32     | kFLOAT       |FLOAT         |kFloat   |FP32     |float32            |
+|TYPE_FP64     |              |DOUBLE        |kDouble  |FP64     |float64            |
+|TYPE_STRING   |              |STRING        |         |BYTES    |dtype(object)      |
+|TYPE_BF16     | kBF16        |BFLOAT16      |         |BF16     |ml_dtypes.bfloat16 |
 
 For TensorRT each value is in the nvinfer1::DataType namespace.
 For example, nvinfer1::DataType::kFLOAT is the 32-bit floating-point datatype.

qa/L0_backend_identity/identity_test.py

@@ -1,6 +1,6 @@
 #!/usr/bin/python
 
-# Copyright 2019-2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -30,6 +30,7 @@
 import sys
 from builtins import range
 
+import ml_dtypes
 import numpy as np
 import requests as httpreq
 import tritonclient.grpc as grpcclient
@@ -201,8 +202,8 @@
             ("identity_nobatch_int8", np.int8, [0]),
             ("identity_nobatch_int8", np.int8, [7]),
             ("identity_bytes", object, [1, 1]),
-            ("identity_bf16", np.float32, [1, 0]),
-            ("identity_bf16", np.float32, [1, 5])
+            ("identity_bf16", ml_dtypes.bfloat16, [1, 0]),
+            ("identity_bf16", ml_dtypes.bfloat16, [1, 5])
         ):
             # yapf: enable
             if np_dtype != object:
@@ -211,10 +212,7 @@
                 in0 = 16384 * np.ones(shape, dtype="int")
                 in0n = np.array([str(x) for x in in0.reshape(in0.size)], dtype=object)
                 input_data = in0n.reshape(in0.shape)
-            if model_name != "identity_bf16":
-                triton_type = np_to_triton_dtype(input_data.dtype)
-            else:
-                triton_type = "BF16"
+            triton_type = np_to_triton_dtype(input_data.dtype)
             inputs = [client_util.InferInput("INPUT0", input_data.shape, triton_type)]
             inputs[0].set_data_from_numpy(input_data)
 

qa/L0_backend_onnxruntime/bfloat16_test.py

@@ -0,0 +1,108 @@
+#!/usr/bin/env python
+# Copyright 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions
+# are met:
+#  * Redistributions of source code must retain the above copyright
+#    notice, this list of conditions and the following disclaimer.
+#  * Redistributions in binary form must reproduce the above copyright
+#    notice, this list of conditions and the following disclaimer in the
+#    documentation and/or other materials provided with the distribution.
+#  * Neither the name of NVIDIA CORPORATION nor the names of its
+#    contributors may be used to endorse or promote products derived
+#    from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import os
+import sys
+import unittest
+
+import ml_dtypes
+import numpy as np
+import pytest
+import tritonclient.grpc as grpcclient
+import tritonclient.http as httpclient
+
+# Client type can be passed as first arg (e.g. python bfloat16_test.py http) or via CLIENT_TYPE env.
+if len(sys.argv) >= 2 and sys.argv[1] in ("http", "grpc"):
+    os.environ["CLIENT_TYPE"] = sys.argv[1]
+    del sys.argv[1]
+
+
+class BFloat16Test(unittest.TestCase):
+    def setUp(self):
+        self.protocol = os.environ.get("CLIENT_TYPE", "http")
+        if self.protocol == "http":
+            self.client_ = httpclient.InferenceServerClient("localhost:8000")
+        else:
+            self.client_ = grpcclient.InferenceServerClient("localhost:8001")
+        self.model_name_ = "add_bf16"
+
+    def _assert_allclose_bf16(self, actual, desired, **kwargs):
+        """Compare bfloat16 arrays by converting to float32 for the check.
+
+        We cannot use np.testing.assert_allclose(actual, desired) directly:
+        isclose() does result_type(y, 1.) and raises DTypePromotionError for
+        bfloat16 in NumPy 2. The error message is misleading—it says
+        "Float16DType and bfloat16" even when both arrays are bfloat16; the
+        real conflict is bfloat16 vs the scalar 1.0 (float64) used inside
+        isclose. Converting to float32 only for the comparison avoids this.
+        """
+        np.testing.assert_allclose(
+            np.asarray(actual, dtype=np.float32),
+            np.asarray(desired, dtype=np.float32),
+            **kwargs,
+        )
+
+    def _infer_bf16(self, input0_data, input1_data):
+        """Helper to run BF16 inference and return the output numpy array."""
+        if self.protocol == "http":
+            input0 = httpclient.InferInput("INPUT0", [5, 5], "BF16")
+            input1 = httpclient.InferInput("INPUT1", [5, 5], "BF16")
+        else:
+            input0 = grpcclient.InferInput("INPUT0", [5, 5], "BF16")
+            input1 = grpcclient.InferInput("INPUT1", [5, 5], "BF16")
+        input0.set_data_from_numpy(input0_data)
+        input1.set_data_from_numpy(input1_data)
+
+        results = self.client_.infer(self.model_name_, [input0, input1])
+        return results.as_numpy("OUTPUT")
+
+    @pytest.mark.parametrize(
+        "input0_val,input1_val,expected_val",
+        [
+            (0.0, 0.0, 0.0),  # zeros
+            (-1.5, 3.5, 2.0),  # negatives / mixed
+            (100.0, 200.0, 300.0),  # large
+            (1e-2, 1e-2, 2e-2),  # small (near underflow)
+            (1.0, -1.0, 0.0),  # cancellation
+            (2.0, 2.0, 4.0),  # identical inputs
+        ],
+    )
+    def test_bf16_add_variants(self, input0_val, input1_val, expected_val):
+        """Run BF16 add for one case: zeros, negatives, large, small, cancellation, or identical."""
+        shape = (5, 5)
+        output = self._infer_bf16(
+            np.full(shape, input0_val, dtype=ml_dtypes.bfloat16),
+            np.full(shape, input1_val, dtype=ml_dtypes.bfloat16),
+        )
+        self.assertEqual(output.dtype, ml_dtypes.bfloat16)
+        self._assert_allclose_bf16(
+            output, np.full(shape, expected_val, dtype=ml_dtypes.bfloat16)
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()

qa/L0_backend_onnxruntime/models/add_bf16/1/model.onnx

@@ -0,0 +1,16 @@
+triton:w
+
+INPUT0
+INPUT1OUTPUT"Addbf16_addZ
+INPUT0
+
+
+Z
+INPUT1
+
+
+b
+OUTPUT
+
+
+B

qa/L0_backend_onnxruntime/models/add_bf16/config.pbtxt

@@ -0,0 +1,50 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions
+# are met:
+#  * Redistributions of source code must retain the above copyright
+#    notice, this list of conditions and the following disclaimer.
+#  * Redistributions in binary form must reproduce the above copyright
+#    notice, this list of conditions and the following disclaimer in the
+#    documentation and/or other materials provided with the distribution.
+#  * Neither the name of NVIDIA CORPORATION nor the names of its
+#    contributors may be used to endorse or promote products derived
+#    from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+platform: "onnxruntime_onnx"
+max_batch_size: 0
+input [
+  {
+    name: "INPUT0"
+    data_type: TYPE_BF16
+    dims: [5, 5]
+  },
+  {
+    name: "INPUT1"
+    data_type: TYPE_BF16
+    dims: [5, 5]
+  }
+]
+output [
+  {
+    name: "OUTPUT"
+    data_type: TYPE_BF16
+    dims: [5, 5]
+  }
+]
+instance_group: {
+  kind: KIND_GPU
+}

qa/L0_backend_onnxruntime/test.sh

@@ -0,0 +1,71 @@
+#!/bin/bash
+# Copyright 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions
+# are met:
+#  * Redistributions of source code must retain the above copyright
+#    notice, this list of conditions and the following disclaimer.
+#  * Redistributions in binary form must reproduce the above copyright
+#    notice, this list of conditions and the following disclaimer in the
+#    documentation and/or other materials provided with the distribution.
+#  * Neither the name of NVIDIA CORPORATION nor the names of its
+#    contributors may be used to endorse or promote products derived
+#    from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+export CUDA_VISIBLE_DEVICES=0
+
+SERVER=/opt/tritonserver/bin/tritonserver
+SERVER_LOG="./inference_server.log"
+CLIENT_LOG="./test.log"
+source ../common/util.sh
+
+rm -f *.log
+
+# BFLOAT16 test
+SERVER_ARGS="--model-repository=`pwd`/models"
+run_server
+if [ "$SERVER_PID" == "0" ]; then
+    echo -e "\n***\n*** Failed to start $SERVER\n***"
+    cat $SERVER_LOG
+    exit 1
+fi
+
+RET=0
+
+set +e
+
+for client_type in http grpc; do
+    CLIENT_LOG="./bfloat16_test_${client_type}.log"
+    python bfloat16_test.py $client_type >>$CLIENT_LOG 2>&1
+    if [ $? -ne 0 ]; then
+        cat $CLIENT_LOG
+        echo -e "\n***\n*** Test Failed ($client_type)\n***"
+        RET=1
+    fi
+done
+
+set -e
+
+kill $SERVER_PID
+wait $SERVER_PID
+
+if [ $RET -eq 0 ]; then
+    echo -e "\n***\n*** Test Passed\n***"
+else
+    echo -e "\n***\n*** Test FAILED\n***"
+fi
+
+exit $RET

qa/L0_backend_python/python_test.py

@@ -1,6 +1,6 @@
 #!/usr/bin/python
 
-# Copyright 2019-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -33,6 +33,7 @@
 import os
 import unittest
 
+import ml_dtypes
 import numpy as np
 import requests as httpreq
 import shm_util
@@ -374,7 +375,7 @@ def test_bf16(self):
             ) as client:
                 # NOTE: Client will truncate FP32 to BF16 internally
                 # since numpy has no built-in BF16 representation.
-                np_input = np.ones(shape, dtype=np.float32)
+                np_input = np.ones(shape, dtype=ml_dtypes.bfloat16)
                 inputs = [
                     httpclient.InferInput(
                         "INPUT0", np_input.shape, "BF16"
@@ -391,8 +392,8 @@ def test_bf16(self):
                 np_output = result.as_numpy("OUTPUT0")
                 self.assertIsNotNone(np_output)
                 # BF16 tensors are held in FP32 when converted to numpy due to
-                # lack of native BF16 support in numpy, so verify that.
-                self.assertEqual(np_output.dtype, np.float32)
+                # lack of native   support in numpy, so verify that.
+                self.assertEqual(np_output.dtype, ml_dtypes.bfloat16)
                 self.assertTrue(np.allclose(np_output, np_input))
 
     def test_infer_pytorch(self):

qa/L0_infer/infer_test.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-# Copyright 2018-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# Copyright 2018-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions
@@ -157,9 +157,6 @@ def _infer_exact_helper(
                     input_dtype,
                     output0_dtype,
                     output1_dtype,
-                    (input_size,),
-                    (input_size,),
-                    (input_size,),
                 ):
                     ensemble_prefix.append(prefix)