From 126fd7e07d6d0fd677a47c7c8e07efdfd4448917 Mon Sep 17 00:00:00 2001 From: Yingge He Date: Fri, 13 Feb 2026 19:15:33 -0800 Subject: [PATCH 01/15] Add BF16 tests for ORT backend --- qa/L0_backend_onnxruntime/bfloat16_test.py | 144 ++++++++++++++++++ .../models/add_bf16/1/model.onnx | 16 ++ .../models/add_bf16/config.pbtxt | 50 ++++++ qa/L0_backend_onnxruntime/test.sh | 71 +++++++++ 4 files changed, 281 insertions(+) create mode 100755 qa/L0_backend_onnxruntime/bfloat16_test.py create mode 100755 qa/L0_backend_onnxruntime/models/add_bf16/1/model.onnx create mode 100644 qa/L0_backend_onnxruntime/models/add_bf16/config.pbtxt create mode 100755 qa/L0_backend_onnxruntime/test.sh diff --git a/qa/L0_backend_onnxruntime/bfloat16_test.py b/qa/L0_backend_onnxruntime/bfloat16_test.py new file mode 100755 index 0000000000..72b5e46f8e --- /dev/null +++ b/qa/L0_backend_onnxruntime/bfloat16_test.py @@ -0,0 +1,144 @@ +#!/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 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") + + def test_bf16_add_variants(self): + """Run multiple BF16 add cases in one test: zeros, negatives, large, small, cancellation, identical.""" + shape = (5, 5) + + # Zeros: 0.0 + 0.0 = 0.0 + output = self._infer_bf16( + np.zeros(shape, dtype=ml_dtypes.bfloat16), + np.zeros(shape, dtype=ml_dtypes.bfloat16), + ) + self.assertEqual(output.dtype, ml_dtypes.bfloat16) + self._assert_allclose_bf16(output, np.zeros(shape, dtype=ml_dtypes.bfloat16)) + + # Negative and mixed: -1.5 + 3.5 = 2.0 + output = self._infer_bf16( + np.full(shape, -1.5, dtype=ml_dtypes.bfloat16), + np.full(shape, 3.5, dtype=ml_dtypes.bfloat16), + ) + self.assertEqual(output.dtype, ml_dtypes.bfloat16) + self._assert_allclose_bf16( + output, np.full(shape, 2.0, dtype=ml_dtypes.bfloat16) + ) + + # Large values within BF16 range: 100 + 200 = 300 + output = self._infer_bf16( + np.full(shape, 100.0, dtype=ml_dtypes.bfloat16), + np.full(shape, 200.0, dtype=ml_dtypes.bfloat16), + ) + self.assertEqual(output.dtype, ml_dtypes.bfloat16) + self._assert_allclose_bf16( + output, np.full(shape, 300.0, dtype=ml_dtypes.bfloat16) + ) + + # Small values (near underflow / precision limit): 0.01 + 0.01 = 0.02 + output = self._infer_bf16( + np.full(shape, 1e-2, dtype=ml_dtypes.bfloat16), + np.full(shape, 1e-2, dtype=ml_dtypes.bfloat16), + ) + self.assertEqual(output.dtype, ml_dtypes.bfloat16) + self._assert_allclose_bf16( + output, np.full(shape, 2e-2, dtype=ml_dtypes.bfloat16) + ) + + # Exact cancellation: 1.0 + (-1.0) = 0.0 + output = self._infer_bf16( + np.full(shape, 1.0, dtype=ml_dtypes.bfloat16), + np.full(shape, -1.0, dtype=ml_dtypes.bfloat16), + ) + self.assertEqual(output.dtype, ml_dtypes.bfloat16) + self._assert_allclose_bf16(output, np.zeros(shape, dtype=ml_dtypes.bfloat16)) + + # Identical inputs: 2.0 + 2.0 = 4.0 + output = self._infer_bf16( + np.full(shape, 2.0, dtype=ml_dtypes.bfloat16), + np.full(shape, 2.0, dtype=ml_dtypes.bfloat16), + ) + self.assertEqual(output.dtype, ml_dtypes.bfloat16) + self._assert_allclose_bf16( + output, np.full(shape, 4.0, dtype=ml_dtypes.bfloat16) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/qa/L0_backend_onnxruntime/models/add_bf16/1/model.onnx b/qa/L0_backend_onnxruntime/models/add_bf16/1/model.onnx new file mode 100755 index 0000000000..c4db6d4232 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/qa/L0_backend_onnxruntime/models/add_bf16/config.pbtxt b/qa/L0_backend_onnxruntime/models/add_bf16/config.pbtxt new file mode 100644 index 0000000000..8e25001a19 --- /dev/null +++ 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 +} \ No newline at end of file diff --git a/qa/L0_backend_onnxruntime/test.sh b/qa/L0_backend_onnxruntime/test.sh new file mode 100755 index 0000000000..74b3b6e168 --- /dev/null +++ b/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 From d92678dcff9fc9d311caea9cacd6ae3f2e2f814f Mon Sep 17 00:00:00 2001 From: Yingge He Date: Fri, 20 Feb 2026 19:39:03 -0800 Subject: [PATCH 02/15] Update related tests and model generations --- docs/user_guide/model_configuration.md | 32 +++++----- qa/L0_backend_identity/identity_test.py | 12 ++-- qa/L0_backend_onnxruntime/bfloat16_test.py | 70 +++++---------------- qa/L0_backend_python/python_test.py | 9 +-- qa/L0_infer/infer_test.py | 5 +- qa/L0_infer_reshape/infer_reshape_test.py | 5 +- qa/L0_infer_variable/infer_variable_test.py | 5 +- qa/L0_infer_zero/infer_zero_test.py | 6 +- qa/L0_trt_bf16_dtype/trt_bf16_dtype_test.py | 7 ++- qa/common/gen_common.py | 13 ++-- qa/common/gen_qa_identity_models.py | 10 +-- qa/common/gen_qa_model_repository | 7 ++- qa/common/gen_qa_models.py | 16 ++--- qa/common/gen_qa_reshape_models.py | 8 +-- qa/common/test_util.py | 14 ++--- 15 files changed, 78 insertions(+), 141 deletions(-) diff --git a/docs/user_guide/model_configuration.md b/docs/user_guide/model_configuration.md index 0c8e0c8875..0643535788 100644 --- a/docs/user_guide/model_configuration.md +++ b/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. diff --git a/qa/L0_backend_identity/identity_test.py b/qa/L0_backend_identity/identity_test.py index a607e4189b..c6a107718d 100755 --- a/qa/L0_backend_identity/identity_test.py +++ b/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) diff --git a/qa/L0_backend_onnxruntime/bfloat16_test.py b/qa/L0_backend_onnxruntime/bfloat16_test.py index 72b5e46f8e..fda724d8e9 100755 --- a/qa/L0_backend_onnxruntime/bfloat16_test.py +++ b/qa/L0_backend_onnxruntime/bfloat16_test.py @@ -31,6 +31,7 @@ import ml_dtypes import numpy as np +import pytest import tritonclient.grpc as grpcclient import tritonclient.http as httpclient @@ -79,64 +80,27 @@ def _infer_bf16(self, input0_data, input1_data): results = self.client_.infer(self.model_name_, [input0, input1]) return results.as_numpy("OUTPUT") - def test_bf16_add_variants(self): - """Run multiple BF16 add cases in one test: zeros, negatives, large, small, cancellation, identical.""" + @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) - - # Zeros: 0.0 + 0.0 = 0.0 - output = self._infer_bf16( - np.zeros(shape, dtype=ml_dtypes.bfloat16), - np.zeros(shape, dtype=ml_dtypes.bfloat16), - ) - self.assertEqual(output.dtype, ml_dtypes.bfloat16) - self._assert_allclose_bf16(output, np.zeros(shape, dtype=ml_dtypes.bfloat16)) - - # Negative and mixed: -1.5 + 3.5 = 2.0 - output = self._infer_bf16( - np.full(shape, -1.5, dtype=ml_dtypes.bfloat16), - np.full(shape, 3.5, dtype=ml_dtypes.bfloat16), - ) - self.assertEqual(output.dtype, ml_dtypes.bfloat16) - self._assert_allclose_bf16( - output, np.full(shape, 2.0, dtype=ml_dtypes.bfloat16) - ) - - # Large values within BF16 range: 100 + 200 = 300 - output = self._infer_bf16( - np.full(shape, 100.0, dtype=ml_dtypes.bfloat16), - np.full(shape, 200.0, dtype=ml_dtypes.bfloat16), - ) - self.assertEqual(output.dtype, ml_dtypes.bfloat16) - self._assert_allclose_bf16( - output, np.full(shape, 300.0, dtype=ml_dtypes.bfloat16) - ) - - # Small values (near underflow / precision limit): 0.01 + 0.01 = 0.02 - output = self._infer_bf16( - np.full(shape, 1e-2, dtype=ml_dtypes.bfloat16), - np.full(shape, 1e-2, dtype=ml_dtypes.bfloat16), - ) - self.assertEqual(output.dtype, ml_dtypes.bfloat16) - self._assert_allclose_bf16( - output, np.full(shape, 2e-2, dtype=ml_dtypes.bfloat16) - ) - - # Exact cancellation: 1.0 + (-1.0) = 0.0 - output = self._infer_bf16( - np.full(shape, 1.0, dtype=ml_dtypes.bfloat16), - np.full(shape, -1.0, dtype=ml_dtypes.bfloat16), - ) - self.assertEqual(output.dtype, ml_dtypes.bfloat16) - self._assert_allclose_bf16(output, np.zeros(shape, dtype=ml_dtypes.bfloat16)) - - # Identical inputs: 2.0 + 2.0 = 4.0 output = self._infer_bf16( - np.full(shape, 2.0, dtype=ml_dtypes.bfloat16), - np.full(shape, 2.0, dtype=ml_dtypes.bfloat16), + 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, 4.0, dtype=ml_dtypes.bfloat16) + output, np.full(shape, expected_val, dtype=ml_dtypes.bfloat16) ) diff --git a/qa/L0_backend_python/python_test.py b/qa/L0_backend_python/python_test.py index c2512600e2..254e74c609 100755 --- a/qa/L0_backend_python/python_test.py +++ b/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): diff --git a/qa/L0_infer/infer_test.py b/qa/L0_infer/infer_test.py index 3d01160b2e..b779032a36 100755 --- a/qa/L0_infer/infer_test.py +++ b/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) diff --git a/qa/L0_infer_reshape/infer_reshape_test.py b/qa/L0_infer_reshape/infer_reshape_test.py index 5445f5d6c9..e55001d7ab 100755 --- a/qa/L0_infer_reshape/infer_reshape_test.py +++ b/qa/L0_infer_reshape/infer_reshape_test.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -# Copyright 2019-2025, 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 @@ -155,9 +155,6 @@ def _full_reshape(self, dtype, input_shapes, output_shapes=None, no_batch=True): dtype, dtype, dtype, - input_shapes[0], - input_shapes[0], - input_shapes[0], ): # model that supports batching for bs in (1, 8): diff --git a/qa/L0_infer_variable/infer_variable_test.py b/qa/L0_infer_variable/infer_variable_test.py index 55a2c7a084..66ecdf60c1 100755 --- a/qa/L0_infer_variable/infer_variable_test.py +++ b/qa/L0_infer_variable/infer_variable_test.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -# Copyright 2019-2025, 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 @@ -140,9 +140,6 @@ def _infer_exact_helper( input_dtype, output0_dtype, output1_dtype, - input_shape, - input_shape, - input_shape, ): ensemble_prefix.append(prefix) diff --git a/qa/L0_infer_zero/infer_zero_test.py b/qa/L0_infer_zero/infer_zero_test.py index d05b383cff..89c77f2554 100755 --- a/qa/L0_infer_zero/infer_zero_test.py +++ b/qa/L0_infer_zero/infer_zero_test.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -# Copyright 2019-2025, 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 @@ -93,9 +93,7 @@ def _full_zero(self, dtype, shapes): ) for name in ["simple_zero", "sequence_zero", "fan_zero"]: - if tu.validate_for_ensemble_model( - name, dtype, dtype, dtype, shapes[0], shapes[0], shapes[0] - ): + if tu.validate_for_ensemble_model(name, dtype, dtype, dtype): # model that supports batching for bs in (1, 8): batch_shapes = [ diff --git a/qa/L0_trt_bf16_dtype/trt_bf16_dtype_test.py b/qa/L0_trt_bf16_dtype/trt_bf16_dtype_test.py index 265c1930b0..50f334516e 100755 --- a/qa/L0_trt_bf16_dtype/trt_bf16_dtype_test.py +++ b/qa/L0_trt_bf16_dtype/trt_bf16_dtype_test.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -# Copyright 2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2024-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 @@ -32,6 +32,7 @@ import unittest +import ml_dtypes import numpy as np import test_util as tu import tritonclient.http as client @@ -49,8 +50,8 @@ def _infer_helper(self, model_name, shape): inputs.append(client.InferInput("INPUT0", shape, "BF16")) inputs.append(client.InferInput("INPUT1", shape, "BF16")) - input0_data = np.ones(shape=shape).astype(np.float32) - input1_data = np.ones(shape=shape).astype(np.float32) + input0_data = np.ones(shape=shape).astype(ml_dtypes.bfloat16) + input1_data = np.ones(shape=shape).astype(ml_dtypes.bfloat16) inputs[0].set_data_from_numpy(input0_data, binary_data=True) inputs[1].set_data_from_numpy(input1_data, binary_data=True) diff --git a/qa/common/gen_common.py b/qa/common/gen_common.py index d53702d604..ca3803b0ed 100644 --- a/qa/common/gen_common.py +++ b/qa/common/gen_common.py @@ -1,4 +1,4 @@ -# Copyright 2023-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2023-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 @@ -28,14 +28,11 @@ from typing import List # Common utilities for model generation scripts +import ml_dtypes import numpy as np np_dtype_string = np.dtype(object) -# Numpy does not support the BF16 datatype natively. -# We use this dummy dtype as a representative for BF16. -np_dtype_bfloat16 = np.dtype([("bf16", object)]) - def np_to_onnx_dtype(np_dtype): import onnx @@ -62,6 +59,8 @@ def np_to_onnx_dtype(np_dtype): return onnx.TensorProto.DOUBLE elif np_dtype == np_dtype_string: return onnx.TensorProto.STRING + elif np_dtype == ml_dtypes.bfloat16: + return onnx.TensorProto.BFLOAT16 return None @@ -88,7 +87,7 @@ def np_to_model_dtype(np_dtype): return "TYPE_FP64" elif np_dtype == np_dtype_string: return "TYPE_STRING" - elif np_dtype == np_dtype_bfloat16: + elif np_dtype == ml_dtypes.bfloat16: return "TYPE_BF16" return None @@ -110,7 +109,7 @@ def np_to_trt_dtype(np_dtype): return trt.float16 elif np_dtype == np.float32: return trt.float32 - elif np_dtype == np_dtype_bfloat16: + elif np_dtype == ml_dtypes.bfloat16: return trt.bfloat16 return None diff --git a/qa/common/gen_qa_identity_models.py b/qa/common/gen_qa_identity_models.py index 7b513d3fbf..17faca8609 100755 --- a/qa/common/gen_qa_identity_models.py +++ b/qa/common/gen_qa_identity_models.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -# Copyright 2019-2025, 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 @@ -47,9 +47,7 @@ def create_ensemble_modelfile( create_savedmodel, models_dir, model_version, io_cnt, max_batch, dtype, shape ): - if not tu.validate_for_ensemble_model( - "zero", dtype, dtype, dtype, shape, shape, shape - ): + if not tu.validate_for_ensemble_model("zero", dtype, dtype, dtype): return emu.create_identity_ensemble_modelfile( @@ -66,9 +64,7 @@ def create_ensemble_modelfile( def create_ensemble_modelconfig( create_savedmodel, models_dir, model_version, io_cnt, max_batch, dtype, shape ): - if not tu.validate_for_ensemble_model( - "zero", dtype, dtype, dtype, shape, shape, shape - ): + if not tu.validate_for_ensemble_model("zero", dtype, dtype, dtype): return emu.create_identity_ensemble_modelconfig( diff --git a/qa/common/gen_qa_model_repository b/qa/common/gen_qa_model_repository index 429d15d736..d2287f3946 100755 --- a/qa/common/gen_qa_model_repository +++ b/qa/common/gen_qa_model_repository @@ -1,5 +1,5 @@ #!/bin/bash -# 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 @@ -153,7 +153,7 @@ apt-get update && \ ln -s /usr/bin/python3 /usr/bin/python -pip3 install "numpy<=1.23.5" setuptools +pip3 install ml_dtypes "numpy<=1.23.5" setuptools pip3 install openvino==$OPENVINO_VERSION @@ -195,7 +195,7 @@ apt-get update && \ protobuf-compiler python3 python3-dev python3-pip ln -s /usr/bin/python3 /usr/bin/python -pip3 install "protobuf<=3.20.1" "numpy<=1.23.5" # TODO: Remove current line DLIS-3838 +pip3 install ml_dtypes pip3 install --upgrade onnx==${ONNX_VERSION} python3 $TRITON_MDLS_SRC_DIR/gen_qa_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$TRITON_MDLS_QA_MODEL @@ -302,6 +302,7 @@ EOF umask 0000 nvidia-smi --query-gpu=compute_cap,compute_mode,driver_version,name,index --format=csv || true nvidia-smi || true +pip3 install ml_dtypes set -e set -x dpkg -l | grep TensorRT diff --git a/qa/common/gen_qa_models.py b/qa/common/gen_qa_models.py index cfce75be39..1fdaafb31f 100755 --- a/qa/common/gen_qa_models.py +++ b/qa/common/gen_qa_models.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 @@ -31,9 +31,9 @@ from builtins import range import gen_ensemble_model_utils as emu +import ml_dtypes import numpy as np from gen_common import ( - np_dtype_bfloat16, np_to_model_dtype, np_to_onnx_dtype, np_to_torch_dtype, @@ -2028,9 +2028,9 @@ def create_fixed_models( if tu.check_gpus_compute_capability(min_capability=8.0): create_fixed_models( FLAGS.models_dir, - np_dtype_bfloat16, - np_dtype_bfloat16, - np_dtype_bfloat16, + ml_dtypes.bfloat16, + ml_dtypes.bfloat16, + ml_dtypes.bfloat16, ) else: print( @@ -2317,9 +2317,9 @@ def create_fixed_models( if tu.check_gpus_compute_capability(min_capability=8.0): create_models( FLAGS.models_dir, - np_dtype_bfloat16, - np_dtype_bfloat16, - np_dtype_bfloat16, + ml_dtypes.bfloat16, + ml_dtypes.bfloat16, + ml_dtypes.bfloat16, (-1, -1), (-1, -1), (-1, -1), diff --git a/qa/common/gen_qa_reshape_models.py b/qa/common/gen_qa_reshape_models.py index d70333c925..5819dcee89 100755 --- a/qa/common/gen_qa_reshape_models.py +++ b/qa/common/gen_qa_reshape_models.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -# Copyright 2019-2025, 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 @@ -522,9 +522,6 @@ def create_ensemble_modelfile( dtype, dtype, dtype, - input_shapes[0], - input_shapes[0], - input_shapes[0], ): return @@ -557,9 +554,6 @@ def create_ensemble_modelconfig( dtype, dtype, dtype, - input_shapes[0], - input_shapes[0], - input_shapes[0], ): return diff --git a/qa/common/test_util.py b/qa/common/test_util.py index ec8108072f..28649e882d 100755 --- a/qa/common/test_util.py +++ b/qa/common/test_util.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 @@ -29,14 +29,11 @@ import json import unittest +import ml_dtypes import numpy as np _last_request_id = 0 -# Numpy does not support the BF16 datatype natively. -# We use this dummy dtype as a representative for BF16. -np_dtype_bfloat16 = np.dtype([("bf16", object)]) - def shape_element_count(shape): cnt = 0 @@ -85,7 +82,7 @@ def validate_for_trt_model( np.uint8, np.float16, np.float32, - np_dtype_bfloat16, + ml_dtypes.bfloat16, ] # FIXME: Remove this check when jetson supports TRT 8.5 (DLIS-4256) if not support_trt_uint8(): @@ -113,9 +110,6 @@ def validate_for_ensemble_model( input_dtype, output0_dtype, output1_dtype, - input_shape, - output0_shape, - output1_shape, ): """Return True if input and output dtypes are supported by the ensemble type.""" @@ -259,7 +253,7 @@ def validate_for_openvino_model( def get_dtype_name(dtype): - if dtype == np_dtype_bfloat16: + if dtype == ml_dtypes.bfloat16: return "bf16" else: return np.dtype(dtype).name From 570bc48e83a88422ee103524786609d5d5f64b61 Mon Sep 17 00:00:00 2001 From: Yingge He Date: Tue, 3 Mar 2026 15:29:17 -0800 Subject: [PATCH 03/15] Fix test --- qa/L0_backend_identity/identity_test.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/qa/L0_backend_identity/identity_test.py b/qa/L0_backend_identity/identity_test.py index c6a107718d..84ae3f6192 100755 --- a/qa/L0_backend_identity/identity_test.py +++ b/qa/L0_backend_identity/identity_test.py @@ -202,6 +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]) ): @@ -212,7 +214,10 @@ 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) - triton_type = np_to_triton_dtype(input_data.dtype) + if model_name == "identity_bf16" and input_data.dtype == np.float32: + triton_type = "BF16" + else: + 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) @@ -232,7 +237,7 @@ print("error: expected 'OUTPUT0'") sys.exit(1) - if model_name == "identity_bf16": + if model_name == "identity_bf16" and input_data.dtype == np.float32: if input_data.shape != output_data.shape: print( "error: expected output shape {} to match input shape {}".format( From c7d325c5083da37330283210aa84ea2abbf2db2c Mon Sep 17 00:00:00 2001 From: Yingge He Date: Tue, 3 Mar 2026 15:44:35 -0800 Subject: [PATCH 04/15] Move bf16 and float32 to different PR --- docs/user_guide/model_configuration.md | 32 ++++++++++----------- qa/L0_backend_identity/identity_test.py | 13 ++++----- qa/L0_backend_python/python_test.py | 7 ++--- qa/L0_trt_bf16_dtype/trt_bf16_dtype_test.py | 5 ++-- qa/common/gen_common.py | 11 +++---- qa/common/gen_qa_model_repository | 7 ++--- qa/common/gen_qa_models.py | 14 ++++----- qa/common/test_util.py | 9 ++++-- 8 files changed, 48 insertions(+), 50 deletions(-) diff --git a/docs/user_guide/model_configuration.md b/docs/user_guide/model_configuration.md index 0643535788..0c8e0c8875 100644 --- a/docs/user_guide/model_configuration.md +++ b/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 |BFLOAT16 | |BF16 |ml_dtypes.bfloat16 | +|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 | | For TensorRT each value is in the nvinfer1::DataType namespace. For example, nvinfer1::DataType::kFLOAT is the 32-bit floating-point datatype. diff --git a/qa/L0_backend_identity/identity_test.py b/qa/L0_backend_identity/identity_test.py index 84ae3f6192..c5b05be473 100755 --- a/qa/L0_backend_identity/identity_test.py +++ b/qa/L0_backend_identity/identity_test.py @@ -30,7 +30,6 @@ import sys from builtins import range -import ml_dtypes import numpy as np import requests as httpreq import tritonclient.grpc as grpcclient @@ -203,9 +202,7 @@ ("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]) + ("identity_bf16", np.float32, [1, 5]) ): # yapf: enable if np_dtype != object: @@ -214,10 +211,10 @@ 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" and input_data.dtype == np.float32: - triton_type = "BF16" - else: + if model_name != "identity_bf16": triton_type = np_to_triton_dtype(input_data.dtype) + else: + triton_type = "BF16" inputs = [client_util.InferInput("INPUT0", input_data.shape, triton_type)] inputs[0].set_data_from_numpy(input_data) @@ -237,7 +234,7 @@ print("error: expected 'OUTPUT0'") sys.exit(1) - if model_name == "identity_bf16" and input_data.dtype == np.float32: + if model_name == "identity_bf16": if input_data.shape != output_data.shape: print( "error: expected output shape {} to match input shape {}".format( diff --git a/qa/L0_backend_python/python_test.py b/qa/L0_backend_python/python_test.py index 254e74c609..3f90af5875 100755 --- a/qa/L0_backend_python/python_test.py +++ b/qa/L0_backend_python/python_test.py @@ -33,7 +33,6 @@ import os import unittest -import ml_dtypes import numpy as np import requests as httpreq import shm_util @@ -375,7 +374,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=ml_dtypes.bfloat16) + np_input = np.ones(shape, dtype=np.float32) inputs = [ httpclient.InferInput( "INPUT0", np_input.shape, "BF16" @@ -392,8 +391,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 support in numpy, so verify that. - self.assertEqual(np_output.dtype, ml_dtypes.bfloat16) + # lack of native BF16 support in numpy, so verify that. + self.assertEqual(np_output.dtype, np.float32) self.assertTrue(np.allclose(np_output, np_input)) def test_infer_pytorch(self): diff --git a/qa/L0_trt_bf16_dtype/trt_bf16_dtype_test.py b/qa/L0_trt_bf16_dtype/trt_bf16_dtype_test.py index 50f334516e..aefc258f94 100755 --- a/qa/L0_trt_bf16_dtype/trt_bf16_dtype_test.py +++ b/qa/L0_trt_bf16_dtype/trt_bf16_dtype_test.py @@ -32,7 +32,6 @@ import unittest -import ml_dtypes import numpy as np import test_util as tu import tritonclient.http as client @@ -50,8 +49,8 @@ def _infer_helper(self, model_name, shape): inputs.append(client.InferInput("INPUT0", shape, "BF16")) inputs.append(client.InferInput("INPUT1", shape, "BF16")) - input0_data = np.ones(shape=shape).astype(ml_dtypes.bfloat16) - input1_data = np.ones(shape=shape).astype(ml_dtypes.bfloat16) + input0_data = np.ones(shape=shape).astype(np.float32) + input1_data = np.ones(shape=shape).astype(np.float32) inputs[0].set_data_from_numpy(input0_data, binary_data=True) inputs[1].set_data_from_numpy(input1_data, binary_data=True) diff --git a/qa/common/gen_common.py b/qa/common/gen_common.py index ca3803b0ed..91f133a0e9 100644 --- a/qa/common/gen_common.py +++ b/qa/common/gen_common.py @@ -28,11 +28,14 @@ from typing import List # Common utilities for model generation scripts -import ml_dtypes import numpy as np np_dtype_string = np.dtype(object) +# Numpy does not support the BF16 datatype natively. +# We use this dummy dtype as a representative for BF16. +np_dtype_bfloat16 = np.dtype([("bf16", object)]) + def np_to_onnx_dtype(np_dtype): import onnx @@ -59,8 +62,6 @@ def np_to_onnx_dtype(np_dtype): return onnx.TensorProto.DOUBLE elif np_dtype == np_dtype_string: return onnx.TensorProto.STRING - elif np_dtype == ml_dtypes.bfloat16: - return onnx.TensorProto.BFLOAT16 return None @@ -87,7 +88,7 @@ def np_to_model_dtype(np_dtype): return "TYPE_FP64" elif np_dtype == np_dtype_string: return "TYPE_STRING" - elif np_dtype == ml_dtypes.bfloat16: + elif np_dtype == np_dtype_bfloat16: return "TYPE_BF16" return None @@ -109,7 +110,7 @@ def np_to_trt_dtype(np_dtype): return trt.float16 elif np_dtype == np.float32: return trt.float32 - elif np_dtype == ml_dtypes.bfloat16: + elif np_dtype == np_dtype_bfloat16: return trt.bfloat16 return None diff --git a/qa/common/gen_qa_model_repository b/qa/common/gen_qa_model_repository index 6eea80fa18..f7cd2b0cbe 100755 --- a/qa/common/gen_qa_model_repository +++ b/qa/common/gen_qa_model_repository @@ -1,5 +1,5 @@ #!/bin/bash -# Copyright 2018-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2018-2025, 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 @@ -153,7 +153,7 @@ apt-get update && \ ln -s /usr/bin/python3 /usr/bin/python -pip3 install ml_dtypes "numpy<=1.23.5" setuptools +pip3 install "numpy<=1.23.5" setuptools pip3 install openvino==$OPENVINO_VERSION @@ -195,7 +195,7 @@ apt-get update && \ protobuf-compiler python3 python3-dev python3-pip ln -s /usr/bin/python3 /usr/bin/python -pip3 install ml_dtypes +pip3 install "protobuf<=3.20.1" "numpy<=1.23.5" # TODO: Remove current line DLIS-3838 pip3 install --upgrade onnx==${ONNX_VERSION} python3 $TRITON_MDLS_SRC_DIR/gen_qa_models.py --onnx --onnx_opset=$ONNX_OPSET --models_dir=$TRITON_MDLS_QA_MODEL @@ -302,7 +302,6 @@ EOF umask 0000 nvidia-smi --query-gpu=compute_cap,compute_mode,driver_version,name,index --format=csv || true nvidia-smi || true -pip3 install ml_dtypes set -e set -x dpkg -l | grep TensorRT diff --git a/qa/common/gen_qa_models.py b/qa/common/gen_qa_models.py index 1fdaafb31f..24227a02fe 100755 --- a/qa/common/gen_qa_models.py +++ b/qa/common/gen_qa_models.py @@ -31,9 +31,9 @@ from builtins import range import gen_ensemble_model_utils as emu -import ml_dtypes import numpy as np from gen_common import ( + np_dtype_bfloat16, np_to_model_dtype, np_to_onnx_dtype, np_to_torch_dtype, @@ -2028,9 +2028,9 @@ def create_fixed_models( if tu.check_gpus_compute_capability(min_capability=8.0): create_fixed_models( FLAGS.models_dir, - ml_dtypes.bfloat16, - ml_dtypes.bfloat16, - ml_dtypes.bfloat16, + np_dtype_bfloat16, + np_dtype_bfloat16, + np_dtype_bfloat16, ) else: print( @@ -2317,9 +2317,9 @@ def create_fixed_models( if tu.check_gpus_compute_capability(min_capability=8.0): create_models( FLAGS.models_dir, - ml_dtypes.bfloat16, - ml_dtypes.bfloat16, - ml_dtypes.bfloat16, + np_dtype_bfloat16, + np_dtype_bfloat16, + np_dtype_bfloat16, (-1, -1), (-1, -1), (-1, -1), diff --git a/qa/common/test_util.py b/qa/common/test_util.py index 28649e882d..b3e66da54d 100755 --- a/qa/common/test_util.py +++ b/qa/common/test_util.py @@ -29,11 +29,14 @@ import json import unittest -import ml_dtypes import numpy as np _last_request_id = 0 +# Numpy does not support the BF16 datatype natively. +# We use this dummy dtype as a representative for BF16. +np_dtype_bfloat16 = np.dtype([("bf16", object)]) + def shape_element_count(shape): cnt = 0 @@ -82,7 +85,7 @@ def validate_for_trt_model( np.uint8, np.float16, np.float32, - ml_dtypes.bfloat16, + np_dtype_bfloat16, ] # FIXME: Remove this check when jetson supports TRT 8.5 (DLIS-4256) if not support_trt_uint8(): @@ -253,7 +256,7 @@ def validate_for_openvino_model( def get_dtype_name(dtype): - if dtype == ml_dtypes.bfloat16: + if dtype == np_dtype_bfloat16: return "bf16" else: return np.dtype(dtype).name From a35923719d9641da7faaacaac8e6d457f958a76d Mon Sep 17 00:00:00 2001 From: Yingge He Date: Tue, 3 Mar 2026 15:48:32 -0800 Subject: [PATCH 05/15] Fix copyrights --- qa/L0_backend_identity/identity_test.py | 2 +- qa/L0_backend_python/python_test.py | 2 +- qa/L0_trt_bf16_dtype/trt_bf16_dtype_test.py | 2 +- qa/common/gen_common.py | 2 +- qa/common/gen_qa_models.py | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/qa/L0_backend_identity/identity_test.py b/qa/L0_backend_identity/identity_test.py index c5b05be473..a607e4189b 100755 --- a/qa/L0_backend_identity/identity_test.py +++ b/qa/L0_backend_identity/identity_test.py @@ -1,6 +1,6 @@ #!/usr/bin/python -# Copyright 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2019-2022, 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 diff --git a/qa/L0_backend_python/python_test.py b/qa/L0_backend_python/python_test.py index 3f90af5875..c2512600e2 100755 --- a/qa/L0_backend_python/python_test.py +++ b/qa/L0_backend_python/python_test.py @@ -1,6 +1,6 @@ #!/usr/bin/python -# Copyright 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2019-2024, 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 diff --git a/qa/L0_trt_bf16_dtype/trt_bf16_dtype_test.py b/qa/L0_trt_bf16_dtype/trt_bf16_dtype_test.py index aefc258f94..265c1930b0 100755 --- a/qa/L0_trt_bf16_dtype/trt_bf16_dtype_test.py +++ b/qa/L0_trt_bf16_dtype/trt_bf16_dtype_test.py @@ -1,5 +1,5 @@ #!/usr/bin/env python3 -# Copyright 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2024, 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 diff --git a/qa/common/gen_common.py b/qa/common/gen_common.py index 91f133a0e9..d53702d604 100644 --- a/qa/common/gen_common.py +++ b/qa/common/gen_common.py @@ -1,4 +1,4 @@ -# Copyright 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2023-2025, 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 diff --git a/qa/common/gen_qa_models.py b/qa/common/gen_qa_models.py index 24227a02fe..cfce75be39 100755 --- a/qa/common/gen_qa_models.py +++ b/qa/common/gen_qa_models.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -# Copyright 2018-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2018-2025, 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 From f400fe5e6c5cdaa35586a590b89c41cf75e8cf9b Mon Sep 17 00:00:00 2001 From: Yingge He Date: Tue, 3 Mar 2026 15:50:17 -0800 Subject: [PATCH 06/15] Add model generate script --- .../models/add_bf16/generate_model.py | 97 +++++++++++++++++++ 1 file changed, 97 insertions(+) create mode 100755 qa/L0_backend_onnxruntime/models/add_bf16/generate_model.py diff --git a/qa/L0_backend_onnxruntime/models/add_bf16/generate_model.py b/qa/L0_backend_onnxruntime/models/add_bf16/generate_model.py new file mode 100755 index 0000000000..97601ae5f5 --- /dev/null +++ b/qa/L0_backend_onnxruntime/models/add_bf16/generate_model.py @@ -0,0 +1,97 @@ +#!/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 onnx + + +def generate_bf16_add_model(models_dir): + """Generate a simple BFLOAT16 Add model (INPUT0 + INPUT1 = OUTPUT).""" + model_name = "add_bf16" + shape = [5, 5] + onnx_dtype = onnx.TensorProto.BFLOAT16 + + add = onnx.helper.make_node("Add", ["INPUT0", "INPUT1"], ["OUTPUT"]) + + input0 = onnx.helper.make_tensor_value_info("INPUT0", onnx_dtype, shape) + input1 = onnx.helper.make_tensor_value_info("INPUT1", onnx_dtype, shape) + output = onnx.helper.make_tensor_value_info("OUTPUT", onnx_dtype, shape) + + graph_proto = onnx.helper.make_graph( + [add], + "bf16_add", + [input0, input1], + [output], + ) + model_def = onnx.helper.make_model(graph_proto, producer_name="triton") + # Cap IR version for older ONNX Runtime (e.g. max supported 11) + model_def.ir_version = min(model_def.ir_version, 11) + # BFLOAT16 support requires opset 13+ + model_def.opset_import[0].version = 13 + + model_dir = os.path.join(models_dir, model_name, "1") + os.makedirs(model_dir, exist_ok=True) + onnx.save(model_def, os.path.join(model_dir, "model.onnx")) + + # Write config.pbtxt + config = """platform: "onnxruntime_onnx" +max_batch_size: 0 +input [ + {{ + name: "INPUT0" + data_type: TYPE_BF16 + dims: {shape} + }}, + {{ + name: "INPUT1" + data_type: TYPE_BF16 + dims: {shape} + }} +] +output [ + {{ + name: "OUTPUT" + data_type: TYPE_BF16 + dims: {shape} + }} +] +""".format( + shape=shape + ) + + config_path = os.path.join(models_dir, model_name, "config.pbtxt") + with open(config_path, "w") as f: + f.write(config) + + print(f"Generated model '{model_name}' in {models_dir}") + + +if __name__ == "__main__": + models_dir = os.path.join(os.getcwd(), "models") + os.makedirs(models_dir, exist_ok=True) + generate_bf16_add_model(models_dir) From ea19f295ba22fdee4f71b5e205b0d306576bb583 Mon Sep 17 00:00:00 2001 From: Yingge He Date: Tue, 3 Mar 2026 15:53:50 -0800 Subject: [PATCH 07/15] Update test --- qa/L0_backend_onnxruntime/bfloat16_test.py | 27 ++++------------------ 1 file changed, 4 insertions(+), 23 deletions(-) diff --git a/qa/L0_backend_onnxruntime/bfloat16_test.py b/qa/L0_backend_onnxruntime/bfloat16_test.py index fda724d8e9..3d13860eb2 100755 --- a/qa/L0_backend_onnxruntime/bfloat16_test.py +++ b/qa/L0_backend_onnxruntime/bfloat16_test.py @@ -29,7 +29,6 @@ import sys import unittest -import ml_dtypes import numpy as np import pytest import tritonclient.grpc as grpcclient @@ -50,22 +49,6 @@ def setUp(self): 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": @@ -95,13 +78,11 @@ 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) + np.full(shape, input0_val, dtype=np.float32), + np.full(shape, input1_val, dtype=np.float32), ) + self.assertEqual(output.dtype, np.float32) + np.testing.assert_allclose(output, expected_val) if __name__ == "__main__": From a5ae694500dab35319c8ac3d8cdffc3c290ceaa3 Mon Sep 17 00:00:00 2001 From: Yingge He Date: Tue, 3 Mar 2026 15:59:29 -0800 Subject: [PATCH 08/15] Update test --- qa/L0_backend_onnxruntime/bfloat16_test.py | 24 +++++++++++----------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/qa/L0_backend_onnxruntime/bfloat16_test.py b/qa/L0_backend_onnxruntime/bfloat16_test.py index 3d13860eb2..85c5b830c3 100755 --- a/qa/L0_backend_onnxruntime/bfloat16_test.py +++ b/qa/L0_backend_onnxruntime/bfloat16_test.py @@ -63,19 +63,19 @@ def _infer_bf16(self, input0_data, 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): + def test_bf16_add_variants(self): """Run BF16 add for one case: zeros, negatives, large, small, cancellation, or identical.""" + 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 + ], + ) + shape = (5, 5) output = self._infer_bf16( np.full(shape, input0_val, dtype=np.float32), From 7e0da9d4a19a2266f5af2a5ade56f5445ee379d2 Mon Sep 17 00:00:00 2001 From: Yingge He Date: Tue, 3 Mar 2026 16:04:34 -0800 Subject: [PATCH 09/15] Fix precommit check --- qa/L0_backend_onnxruntime/bfloat16_test.py | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/qa/L0_backend_onnxruntime/bfloat16_test.py b/qa/L0_backend_onnxruntime/bfloat16_test.py index 85c5b830c3..6d859874bf 100755 --- a/qa/L0_backend_onnxruntime/bfloat16_test.py +++ b/qa/L0_backend_onnxruntime/bfloat16_test.py @@ -30,7 +30,6 @@ import unittest import numpy as np -import pytest import tritonclient.grpc as grpcclient import tritonclient.http as httpclient @@ -65,16 +64,14 @@ def _infer_bf16(self, input0_data, input1_data): def test_bf16_add_variants(self): """Run BF16 add for one case: zeros, negatives, large, small, cancellation, or identical.""" - 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 - ], - ) + 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 + ] shape = (5, 5) output = self._infer_bf16( From 4824162a523417aa1cc512a65d87a77697302e42 Mon Sep 17 00:00:00 2001 From: Yingge He Date: Tue, 3 Mar 2026 16:47:47 -0800 Subject: [PATCH 10/15] Update test --- qa/L0_backend_onnxruntime/bfloat16_test.py | 27 +++++++++++----------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/qa/L0_backend_onnxruntime/bfloat16_test.py b/qa/L0_backend_onnxruntime/bfloat16_test.py index 6d859874bf..aba7e00c02 100755 --- a/qa/L0_backend_onnxruntime/bfloat16_test.py +++ b/qa/L0_backend_onnxruntime/bfloat16_test.py @@ -47,15 +47,16 @@ def setUp(self): else: self.client_ = grpcclient.InferenceServerClient("localhost:8001") self.model_name_ = "add_bf16" + self.shape_ = [5, 5] 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") + input0 = httpclient.InferInput("INPUT0", self.shape_, "BF16") + input1 = httpclient.InferInput("INPUT1", self.shape_, "BF16") else: - input0 = grpcclient.InferInput("INPUT0", [5, 5], "BF16") - input1 = grpcclient.InferInput("INPUT1", [5, 5], "BF16") + input0 = grpcclient.InferInput("INPUT0", self.shape_, "BF16") + input1 = grpcclient.InferInput("INPUT1", self.shape_, "BF16") input0.set_data_from_numpy(input0_data) input1.set_data_from_numpy(input1_data) @@ -64,22 +65,20 @@ def _infer_bf16(self, input0_data, input1_data): def test_bf16_add_variants(self): """Run BF16 add for one case: zeros, negatives, large, small, cancellation, or identical.""" - input0_val, input1_val, expected_val = [ + for input0_val, input1_val, expected_val in [ (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 - ] - - shape = (5, 5) - output = self._infer_bf16( - np.full(shape, input0_val, dtype=np.float32), - np.full(shape, input1_val, dtype=np.float32), - ) - self.assertEqual(output.dtype, np.float32) - np.testing.assert_allclose(output, expected_val) + ]: + output = self._infer_bf16( + np.full(self.shape_, input0_val, dtype=np.float32), + np.full(self.shape_, input1_val, dtype=np.float32), + ) + self.assertEqual(output.dtype, np.float32) + np.testing.assert_allclose(output, expected_val) if __name__ == "__main__": From 7c9cafe95565a863f27c2a5aa6220a03c2907997 Mon Sep 17 00:00:00 2001 From: Yingge He Date: Wed, 4 Mar 2026 12:55:07 -0800 Subject: [PATCH 11/15] Update test --- ...te_model.py => gen_add_bf16_onnx_model.py} | 7 ++- .../models/add_bf16/1/model.onnx | 16 ------ .../models/add_bf16/config.pbtxt | 50 ------------------- .../{bfloat16_test.py => test.py} | 6 ++- qa/L0_backend_onnxruntime/test.sh | 12 ++++- 5 files changed, 19 insertions(+), 72 deletions(-) rename qa/L0_backend_onnxruntime/{models/add_bf16/generate_model.py => gen_add_bf16_onnx_model.py} (94%) delete mode 100755 qa/L0_backend_onnxruntime/models/add_bf16/1/model.onnx delete mode 100644 qa/L0_backend_onnxruntime/models/add_bf16/config.pbtxt rename qa/L0_backend_onnxruntime/{bfloat16_test.py => test.py} (93%) diff --git a/qa/L0_backend_onnxruntime/models/add_bf16/generate_model.py b/qa/L0_backend_onnxruntime/gen_add_bf16_onnx_model.py similarity index 94% rename from qa/L0_backend_onnxruntime/models/add_bf16/generate_model.py rename to qa/L0_backend_onnxruntime/gen_add_bf16_onnx_model.py index 97601ae5f5..87d5a64c0a 100755 --- a/qa/L0_backend_onnxruntime/models/add_bf16/generate_model.py +++ b/qa/L0_backend_onnxruntime/gen_add_bf16_onnx_model.py @@ -25,6 +25,9 @@ # (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +# Generates the add_bf16 ONNX model and Triton config. +# Model: element-wise Add in BFLOAT16 (INPUT0 + INPUT1 = OUTPUT), ONNX Runtime backend. import os import onnx @@ -33,7 +36,7 @@ def generate_bf16_add_model(models_dir): """Generate a simple BFLOAT16 Add model (INPUT0 + INPUT1 = OUTPUT).""" model_name = "add_bf16" - shape = [5, 5] + shape = [1] onnx_dtype = onnx.TensorProto.BFLOAT16 add = onnx.helper.make_node("Add", ["INPUT0", "INPUT1"], ["OUTPUT"]) @@ -44,7 +47,7 @@ def generate_bf16_add_model(models_dir): graph_proto = onnx.helper.make_graph( [add], - "bf16_add", + model_name, [input0, input1], [output], ) diff --git a/qa/L0_backend_onnxruntime/models/add_bf16/1/model.onnx b/qa/L0_backend_onnxruntime/models/add_bf16/1/model.onnx deleted file mode 100755 index c4db6d4232..0000000000 --- a/qa/L0_backend_onnxruntime/models/add_bf16/1/model.onnx +++ /dev/null @@ -1,16 +0,0 @@ - triton:w - -INPUT0 -INPUT1OUTPUT"Addbf16_addZ -INPUT0 -  - -Z -INPUT1 -  - -b -OUTPUT -  - -B \ No newline at end of file diff --git a/qa/L0_backend_onnxruntime/models/add_bf16/config.pbtxt b/qa/L0_backend_onnxruntime/models/add_bf16/config.pbtxt deleted file mode 100644 index 8e25001a19..0000000000 --- a/qa/L0_backend_onnxruntime/models/add_bf16/config.pbtxt +++ /dev/null @@ -1,50 +0,0 @@ -# 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 -} \ No newline at end of file diff --git a/qa/L0_backend_onnxruntime/bfloat16_test.py b/qa/L0_backend_onnxruntime/test.py similarity index 93% rename from qa/L0_backend_onnxruntime/bfloat16_test.py rename to qa/L0_backend_onnxruntime/test.py index aba7e00c02..c09eb8352e 100755 --- a/qa/L0_backend_onnxruntime/bfloat16_test.py +++ b/qa/L0_backend_onnxruntime/test.py @@ -47,7 +47,7 @@ def setUp(self): else: self.client_ = grpcclient.InferenceServerClient("localhost:8001") self.model_name_ = "add_bf16" - self.shape_ = [5, 5] + self.shape_ = [1] def _infer_bf16(self, input0_data, input1_data): """Helper to run BF16 inference and return the output numpy array.""" @@ -78,7 +78,9 @@ def test_bf16_add_variants(self): np.full(self.shape_, input1_val, dtype=np.float32), ) self.assertEqual(output.dtype, np.float32) - np.testing.assert_allclose(output, expected_val) + # TODO: BF16 to FP32 conversion loses precision. Remove rtol and atol in TRI-801. + # BF16 has ~3 decimal digits; use relaxed tol for computed values + np.testing.assert_allclose(output, expected_val, rtol=1e-2, atol=1e-3) if __name__ == "__main__": diff --git a/qa/L0_backend_onnxruntime/test.sh b/qa/L0_backend_onnxruntime/test.sh index 74b3b6e168..c7b1db56ae 100755 --- a/qa/L0_backend_onnxruntime/test.sh +++ b/qa/L0_backend_onnxruntime/test.sh @@ -33,8 +33,16 @@ CLIENT_LOG="./test.log" source ../common/util.sh rm -f *.log +rm -rf models # BFLOAT16 test +# Generate the model +mkdir -p models/add_bf16/1 +set +e +pip install onnx==1.20.1 +python gen_add_bf16_onnx_model.py +set -e + SERVER_ARGS="--model-repository=`pwd`/models" run_server if [ "$SERVER_PID" == "0" ]; then @@ -48,8 +56,8 @@ 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 + CLIENT_LOG="./test_${client_type}.log" + python test.py $client_type >>$CLIENT_LOG 2>&1 if [ $? -ne 0 ]; then cat $CLIENT_LOG echo -e "\n***\n*** Test Failed ($client_type)\n***" From a555d14744d211338d606cc41e02635fd1bab1b2 Mon Sep 17 00:00:00 2001 From: Yingge He Date: Wed, 4 Mar 2026 18:05:54 -0800 Subject: [PATCH 12/15] Update tests --- .../gen_add_bf16_onnx_model.py | 2 +- qa/L0_backend_onnxruntime/test.py | 10 ++-------- qa/L0_backend_onnxruntime/test.sh | 17 +++++++++++++++-- 3 files changed, 18 insertions(+), 11 deletions(-) diff --git a/qa/L0_backend_onnxruntime/gen_add_bf16_onnx_model.py b/qa/L0_backend_onnxruntime/gen_add_bf16_onnx_model.py index 87d5a64c0a..e3fc39b38e 100755 --- a/qa/L0_backend_onnxruntime/gen_add_bf16_onnx_model.py +++ b/qa/L0_backend_onnxruntime/gen_add_bf16_onnx_model.py @@ -1,4 +1,4 @@ -#!/usr/bin/env python +#!/usr/bin/env python3 # Copyright 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Redistribution and use in source and binary forms, with or without diff --git a/qa/L0_backend_onnxruntime/test.py b/qa/L0_backend_onnxruntime/test.py index c09eb8352e..1a5437cdce 100755 --- a/qa/L0_backend_onnxruntime/test.py +++ b/qa/L0_backend_onnxruntime/test.py @@ -1,4 +1,4 @@ -#!/usr/bin/env python +#!/usr/bin/env python3 # Copyright 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Redistribution and use in source and binary forms, with or without @@ -26,18 +26,12 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. import os -import sys import unittest import numpy as np 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): @@ -64,7 +58,7 @@ def _infer_bf16(self, input0_data, input1_data): return results.as_numpy("OUTPUT") def test_bf16_add_variants(self): - """Run BF16 add for one case: zeros, negatives, large, small, cancellation, or identical.""" + """Run BF16 add across multiple cases: zeros, negatives, large, small, cancellation, and identical.""" for input0_val, input1_val, expected_val in [ (0.0, 0.0, 0.0), # zeros (-1.5, 3.5, 2.0), # negatives / mixed diff --git a/qa/L0_backend_onnxruntime/test.sh b/qa/L0_backend_onnxruntime/test.sh index c7b1db56ae..fb4f73f466 100755 --- a/qa/L0_backend_onnxruntime/test.sh +++ b/qa/L0_backend_onnxruntime/test.sh @@ -39,8 +39,19 @@ rm -rf models # Generate the model mkdir -p models/add_bf16/1 set +e + pip install onnx==1.20.1 +if [ $? -ne 0 ]; then + echo -e "\n***\n*** Failed to install onnx dependency\n***" + RET=1 +fi + python gen_add_bf16_onnx_model.py +if [ $? -ne 0 ]; then + echo -e "\n***\n*** Failed to generate BFLOAT16 ONNX model\n***" + RET=1 +fi + set -e SERVER_ARGS="--model-repository=`pwd`/models" @@ -48,7 +59,7 @@ run_server if [ "$SERVER_PID" == "0" ]; then echo -e "\n***\n*** Failed to start $SERVER\n***" cat $SERVER_LOG - exit 1 + RET=1 fi RET=0 @@ -56,14 +67,16 @@ RET=0 set +e for client_type in http grpc; do + export CLIENT_TYPE=$client_type CLIENT_LOG="./test_${client_type}.log" - python test.py $client_type >>$CLIENT_LOG 2>&1 + python test.py >>$CLIENT_LOG 2>&1 if [ $? -ne 0 ]; then cat $CLIENT_LOG echo -e "\n***\n*** Test Failed ($client_type)\n***" RET=1 fi done +unset CLIENT_TYPE set -e From b9d6e30274b9afa36d57fe49b63470ae831c5c09 Mon Sep 17 00:00:00 2001 From: Yingge He Date: Wed, 4 Mar 2026 18:20:37 -0800 Subject: [PATCH 13/15] Fix tests --- qa/L0_backend_onnxruntime/test.sh | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/qa/L0_backend_onnxruntime/test.sh b/qa/L0_backend_onnxruntime/test.sh index fb4f73f466..4cb825ec9b 100755 --- a/qa/L0_backend_onnxruntime/test.sh +++ b/qa/L0_backend_onnxruntime/test.sh @@ -35,6 +35,8 @@ source ../common/util.sh rm -f *.log rm -rf models +RET=0 + # BFLOAT16 test # Generate the model mkdir -p models/add_bf16/1 @@ -59,11 +61,9 @@ run_server if [ "$SERVER_PID" == "0" ]; then echo -e "\n***\n*** Failed to start $SERVER\n***" cat $SERVER_LOG - RET=1 + exit 1 fi -RET=0 - set +e for client_type in http grpc; do From abdb2601a99ecdce0e62ba58b320b566b98f1d0c Mon Sep 17 00:00:00 2001 From: Yingge He Date: Wed, 4 Mar 2026 18:30:05 -0800 Subject: [PATCH 14/15] exit test --- qa/L0_backend_onnxruntime/test.sh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/qa/L0_backend_onnxruntime/test.sh b/qa/L0_backend_onnxruntime/test.sh index 4cb825ec9b..f4e349c2be 100755 --- a/qa/L0_backend_onnxruntime/test.sh +++ b/qa/L0_backend_onnxruntime/test.sh @@ -45,13 +45,13 @@ set +e pip install onnx==1.20.1 if [ $? -ne 0 ]; then echo -e "\n***\n*** Failed to install onnx dependency\n***" - RET=1 + exit 1 fi python gen_add_bf16_onnx_model.py if [ $? -ne 0 ]; then echo -e "\n***\n*** Failed to generate BFLOAT16 ONNX model\n***" - RET=1 + exit 1 fi set -e From 356006c06480259a5e5d3f24d211c27fdbda7dc3 Mon Sep 17 00:00:00 2001 From: Yingge He Date: Thu, 5 Mar 2026 10:23:34 -0800 Subject: [PATCH 15/15] Revert model generation --- qa/L0_infer/infer_test.py | 5 ++++- qa/L0_infer_reshape/infer_reshape_test.py | 5 ++++- qa/L0_infer_variable/infer_variable_test.py | 5 ++++- qa/L0_infer_zero/infer_zero_test.py | 6 ++++-- qa/common/gen_qa_identity_models.py | 10 +++++++--- qa/common/gen_qa_reshape_models.py | 8 +++++++- qa/common/test_util.py | 5 ++++- 7 files changed, 34 insertions(+), 10 deletions(-) diff --git a/qa/L0_infer/infer_test.py b/qa/L0_infer/infer_test.py index b779032a36..3d01160b2e 100755 --- a/qa/L0_infer/infer_test.py +++ b/qa/L0_infer/infer_test.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -# Copyright 2018-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2018-2025, 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,6 +157,9 @@ def _infer_exact_helper( input_dtype, output0_dtype, output1_dtype, + (input_size,), + (input_size,), + (input_size,), ): ensemble_prefix.append(prefix) diff --git a/qa/L0_infer_reshape/infer_reshape_test.py b/qa/L0_infer_reshape/infer_reshape_test.py index e55001d7ab..5445f5d6c9 100755 --- a/qa/L0_infer_reshape/infer_reshape_test.py +++ b/qa/L0_infer_reshape/infer_reshape_test.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -# Copyright 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2019-2025, 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 @@ -155,6 +155,9 @@ def _full_reshape(self, dtype, input_shapes, output_shapes=None, no_batch=True): dtype, dtype, dtype, + input_shapes[0], + input_shapes[0], + input_shapes[0], ): # model that supports batching for bs in (1, 8): diff --git a/qa/L0_infer_variable/infer_variable_test.py b/qa/L0_infer_variable/infer_variable_test.py index 66ecdf60c1..55a2c7a084 100755 --- a/qa/L0_infer_variable/infer_variable_test.py +++ b/qa/L0_infer_variable/infer_variable_test.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -# Copyright 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2019-2025, 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 @@ -140,6 +140,9 @@ def _infer_exact_helper( input_dtype, output0_dtype, output1_dtype, + input_shape, + input_shape, + input_shape, ): ensemble_prefix.append(prefix) diff --git a/qa/L0_infer_zero/infer_zero_test.py b/qa/L0_infer_zero/infer_zero_test.py index 89c77f2554..d05b383cff 100755 --- a/qa/L0_infer_zero/infer_zero_test.py +++ b/qa/L0_infer_zero/infer_zero_test.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -# Copyright 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2019-2025, 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 @@ -93,7 +93,9 @@ def _full_zero(self, dtype, shapes): ) for name in ["simple_zero", "sequence_zero", "fan_zero"]: - if tu.validate_for_ensemble_model(name, dtype, dtype, dtype): + if tu.validate_for_ensemble_model( + name, dtype, dtype, dtype, shapes[0], shapes[0], shapes[0] + ): # model that supports batching for bs in (1, 8): batch_shapes = [ diff --git a/qa/common/gen_qa_identity_models.py b/qa/common/gen_qa_identity_models.py index 17faca8609..7b513d3fbf 100755 --- a/qa/common/gen_qa_identity_models.py +++ b/qa/common/gen_qa_identity_models.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -# Copyright 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2019-2025, 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 @@ -47,7 +47,9 @@ def create_ensemble_modelfile( create_savedmodel, models_dir, model_version, io_cnt, max_batch, dtype, shape ): - if not tu.validate_for_ensemble_model("zero", dtype, dtype, dtype): + if not tu.validate_for_ensemble_model( + "zero", dtype, dtype, dtype, shape, shape, shape + ): return emu.create_identity_ensemble_modelfile( @@ -64,7 +66,9 @@ def create_ensemble_modelfile( def create_ensemble_modelconfig( create_savedmodel, models_dir, model_version, io_cnt, max_batch, dtype, shape ): - if not tu.validate_for_ensemble_model("zero", dtype, dtype, dtype): + if not tu.validate_for_ensemble_model( + "zero", dtype, dtype, dtype, shape, shape, shape + ): return emu.create_identity_ensemble_modelconfig( diff --git a/qa/common/gen_qa_reshape_models.py b/qa/common/gen_qa_reshape_models.py index 5819dcee89..d70333c925 100755 --- a/qa/common/gen_qa_reshape_models.py +++ b/qa/common/gen_qa_reshape_models.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -# Copyright 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2019-2025, 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 @@ -522,6 +522,9 @@ def create_ensemble_modelfile( dtype, dtype, dtype, + input_shapes[0], + input_shapes[0], + input_shapes[0], ): return @@ -554,6 +557,9 @@ def create_ensemble_modelconfig( dtype, dtype, dtype, + input_shapes[0], + input_shapes[0], + input_shapes[0], ): return diff --git a/qa/common/test_util.py b/qa/common/test_util.py index b3e66da54d..ec8108072f 100755 --- a/qa/common/test_util.py +++ b/qa/common/test_util.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 -# Copyright 2018-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright 2018-2025, 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 @@ -113,6 +113,9 @@ def validate_for_ensemble_model( input_dtype, output0_dtype, output1_dtype, + input_shape, + output0_shape, + output1_shape, ): """Return True if input and output dtypes are supported by the ensemble type."""